perf(streaming): stream tool input up to 24x faster by skipping re-parses inside open strings

src/anthropic/lib/streaming/_beta_messages.py

@@ -29,6 +29,7 @@
 )
 from ..._streaming import Stream, AsyncStream
 from ...types.beta import BetaRawMessageStreamEvent
+from ._partial_json import PartialJsonBuffer
 from ..._utils._utils import is_given
 from .._parse._response import ResponseFormatT, parse_text
 from ...types.beta.parsed_beta_message import ParsedBetaMessage, ParsedBetaContentBlock
@@ -54,7 +55,7 @@ def __init__(
         self.text_stream = self.__stream_text__()
         self._iterator = self.__stream__()
         self.__final_message_snapshot: ParsedBetaMessage[ResponseFormatT] | None = None
-        self.__json_bufs: dict[int, bytes] = {}
+        self.__json_bufs: dict[int, PartialJsonBuffer] = {}
         self.__output_format = output_format
         share_message_snapshot(raw_stream, self._current_snapshot_or_none)
 
@@ -220,7 +221,7 @@ def __init__(
         self.text_stream = self.__stream_text__()
         self._iterator = self.__stream__()
         self.__final_message_snapshot: ParsedBetaMessage[ResponseFormatT] | None = None
-        self.__json_bufs: dict[int, bytes] = {}
+        self.__json_bufs: dict[int, PartialJsonBuffer] = {}
         self.__output_format = output_format
         share_message_snapshot(raw_stream, self._current_snapshot_or_none)
 
@@ -481,7 +482,7 @@ def accumulate_event(
     *,
     event: BetaRawMessageStreamEvent,
     current_snapshot: ParsedBetaMessage[ResponseFormatT] | None,
-    json_bufs: dict[int, bytes],
+    json_bufs: dict[int, PartialJsonBuffer],
     request_headers: httpx2.Headers,
     output_format: ResponseFormatT | NotGiven = not_given,
 ) -> ParsedBetaMessage[ResponseFormatT]:
@@ -527,23 +528,21 @@ def accumulate_event(
             if isinstance(content, TRACKS_TOOL_INPUT):
                 from jiter import from_json
 
-                json_buf = json_bufs.get(event.index, b"")
-                json_buf += bytes(event.delta.partial_json, "utf-8")
+                json_buf = json_bufs.setdefault(event.index, PartialJsonBuffer())
+                extends_open_string = json_buf.feed(bytes(event.delta.partial_json, "utf-8"))
 
                 if json_buf:
                     try:
                         anthropic_beta = request_headers.get("anthropic-beta", "") if request_headers else ""
 
                         if "fine-grained-tool-streaming-2025-05-14" in anthropic_beta:
-                            content.input = from_json(json_buf, partial_mode="trailing-strings")
-                        else:
-                            content.input = from_json(json_buf, partial_mode=True)
+                            content.input = from_json(json_buf.getvalue(), partial_mode="trailing-strings")
+                        elif not extends_open_string:
+                            content.input = from_json(json_buf.getvalue(), partial_mode=True)
                     except ValueError as e:
                         raise ValueError(
-                            f"Unable to parse tool parameter JSON from model. Please retry your request or adjust your prompt. Error: {e}. JSON: {json_buf.decode('utf-8')}"
+                            f"Unable to parse tool parameter JSON from model. Please retry your request or adjust your prompt. Error: {e}. JSON: {json_buf.getvalue().decode('utf-8')}"
                         ) from e
-
-                json_bufs[event.index] = json_buf
         elif event.delta.type == "citations_delta":
             if content.type == "text":
                 if not content.citations:

src/anthropic/lib/streaming/_messages.py

@@ -26,6 +26,7 @@
 from ..._models import build, construct_type, construct_type_unchecked
 from .._tracing import share_message_snapshot, withhold_message_snapshot
 from ..._streaming import Stream, AsyncStream
+from ._partial_json import PartialJsonBuffer
 from ..._utils._utils import is_given
 from .._parse._response import ResponseFormatT, parse_text
 from ...types.parsed_message import ParsedMessage, ParsedContentBlock
@@ -51,7 +52,7 @@ def __init__(
         self.text_stream = self.__stream_text__()
         self._iterator = self.__stream__()
         self.__final_message_snapshot: ParsedMessage[ResponseFormatT] | None = None
-        self.__json_bufs: dict[int, bytes] = {}
+        self.__json_bufs: dict[int, PartialJsonBuffer] = {}
         self.__output_format = output_format
         share_message_snapshot(raw_stream, self._current_snapshot_or_none)
 
@@ -216,7 +217,7 @@ def __init__(
         self.text_stream = self.__stream_text__()
         self._iterator = self.__stream__()
         self.__final_message_snapshot: ParsedMessage[ResponseFormatT] | None = None
-        self.__json_bufs: dict[int, bytes] = {}
+        self.__json_bufs: dict[int, PartialJsonBuffer] = {}
         self.__output_format = output_format
         share_message_snapshot(raw_stream, self._current_snapshot_or_none)
 
@@ -465,7 +466,7 @@ def accumulate_event(
     *,
     event: RawMessageStreamEvent,
     current_snapshot: ParsedMessage[ResponseFormatT] | None,
-    json_bufs: dict[int, bytes],
+    json_bufs: dict[int, PartialJsonBuffer],
     output_format: ResponseFormatT | NotGiven = not_given,
 ) -> ParsedMessage[ResponseFormatT]:
     if not isinstance(cast(Any, event), BaseModel):
@@ -502,18 +503,16 @@ def accumulate_event(
             if isinstance(content, TRACKS_TOOL_INPUT):
                 from jiter import from_json
 
-                json_buf = json_bufs.get(event.index, b"")
-                json_buf += bytes(event.delta.partial_json, "utf-8")
+                json_buf = json_bufs.setdefault(event.index, PartialJsonBuffer())
+                extends_open_string = json_buf.feed(bytes(event.delta.partial_json, "utf-8"))
 
-                if json_buf:
+                if json_buf and not extends_open_string:
                     try:
-                        content.input = from_json(json_buf, partial_mode=True)
+                        content.input = from_json(json_buf.getvalue(), partial_mode=True)
                     except ValueError as e:
                         raise ValueError(
-                            f"Unable to parse tool parameter JSON from model. Please retry your request or adjust your prompt. Error: {e}. JSON: {json_buf.decode('utf-8')}"
+                            f"Unable to parse tool parameter JSON from model. Please retry your request or adjust your prompt. Error: {e}. JSON: {json_buf.getvalue().decode('utf-8')}"
                         ) from e
-
-                json_bufs[event.index] = json_buf
         elif event.delta.type == "citations_delta":
             if content.type == "text":
                 if not content.citations:

src/anthropic/lib/streaming/_partial_json.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import re
+
+# Bytes that end a JSON string, start an escape or are not allowed in one.
+_STRING_SPECIAL = re.compile(rb'["\\\x00-\x1f]')
+_SIMPLE_ESCAPES = frozenset(b'"\\/bfnrt')
+_HEX_DIGITS = frozenset(b"0123456789abcdefABCDEF")
+
+
+class PartialJsonBuffer:
+    """Collects the `partial_json` deltas of one content block.
+
+    `from_json(..., partial_mode=True)` drops a trailing unclosed string, so a delta that only extends
+    an open string does not change the parsed value and the buffer does not need to be parsed again.
+    Anything that is not valid inside a JSON string marks the buffer as unsafe to skip from then on,
+    so that the parser still sees it and raises the same error.
+    """
+
+    def __init__(self) -> None:
+        self._buf = bytearray()
+        self._in_string = False
+        self._invalid = False
+        # 0: no escape, 1: after a backslash, 2-5: reading the hex digits of `\u`
+        self._escape = 0
+        self._code = 0
+        self._high_surrogate = False
+
+    def __len__(self) -> int:
+        return len(self._buf)
+
+    def getvalue(self) -> bytes:
+        return bytes(self._buf)
+
+    def feed(self, data: bytes) -> bool:
+        """Appends `data` and returns whether it only extended an open string with valid content."""
+        self._buf += data
+        if self._invalid:
+            return False
+
+        extends_open_string = self._in_string
+        pos = 0
+        end = len(data)
+        while pos < end:
+            if not self._in_string:
+                quote = data.find(b'"', pos)
+                if quote < 0:
+                    break
+                self._in_string = True
+                pos = quote + 1
+                continue
+
+            char = data[pos]
+            if self._escape == 0:
+                if self._high_surrogate and char != 0x5C:  # a high surrogate must be followed by `\u`
+                    self._invalid = True
+                    return False
+                match = _STRING_SPECIAL.search(data, pos)
+                if match is None:
+                    break
+                pos = match.end()
+                char = data[pos - 1]
+                if char == 0x22:  # '"'
+                    self._in_string = False
+                    extends_open_string = False
+                elif char == 0x5C:  # '\\'
+                    self._escape = 1
+                else:
+                    self._invalid = True
+                    return False
+            elif self._escape == 1:
+                pos += 1
+                if char == 0x75:  # 'u'
+                    self._escape = 2
+                    self._code = 0
+                elif char in _SIMPLE_ESCAPES and not self._high_surrogate:
+                    self._escape = 0
+                else:
+                    self._invalid = True
+                    return False
+            else:
+                pos += 1
+                if char not in _HEX_DIGITS:
+                    self._invalid = True
+                    return False
+                self._code = self._code * 16 + int(chr(char), 16)
+                self._escape += 1
+                if self._escape == 6:
+                    self._escape = 0
+                    is_high = 0xD800 <= self._code <= 0xDBFF
+                    is_low = 0xDC00 <= self._code <= 0xDFFF
+                    if self._high_surrogate != is_low:
+                        self._invalid = True
+                        return False
+                    self._high_surrogate = is_high
+
+        return extends_open_string

tests/lib/streaming/helpers.py

@@ -1,9 +1,13 @@
 from __future__ import annotations
 
 import os
-from typing import TypeVar, Iterator
+import json
+from typing import Any, TypeVar, Iterator
 from typing_extensions import AsyncIterator
 
+import jiter
+import pytest
+
 _T = TypeVar("_T")
 
 
@@ -26,3 +30,79 @@ async def to_async_iter(iter: Iterator[_T]) -> AsyncIterator[_T]:
     """Convert a synchronous iterator to an asynchronous one."""
     for event in iter:
         yield event
+
+
+LONG_TOOL_INPUT = json.dumps({"path": "notes.md", "content": 'a "quoted" café line\\ \n' * 2000})
+
+
+def get_long_tool_input_response(tool_input: str = LONG_TOOL_INPUT) -> Iterator[bytes]:
+    """A `tool_use` block whose input holds a long string, streamed in small `input_json_delta` events."""
+    events: list[dict[str, Any]] = [
+        {
+            "type": "message_start",
+            "message": {
+                "id": "msg_01",
+                "type": "message",
+                "role": "assistant",
+                "model": "claude-sonnet-4-20250514",
+                "content": [],
+                "stop_reason": None,
+                "stop_sequence": None,
+                "usage": {"input_tokens": 10, "output_tokens": 1},
+            },
+        },
+        {
+            "type": "content_block_start",
+            "index": 0,
+            "content_block": {
+                "type": "tool_use",
+                "id": "toolu_01",
+                "name": "write_file",
+                "input": {},
+                "caller": {"type": "direct"},
+            },
+        },
+        *(
+            {
+                "type": "content_block_delta",
+                "index": 0,
+                "delta": {"type": "input_json_delta", "partial_json": tool_input[i : i + 8]},
+            }
+            for i in range(0, len(tool_input), 8)
+        ),
+        {"type": "content_block_stop", "index": 0},
+        {
+            "type": "message_delta",
+            "delta": {"stop_reason": "tool_use", "stop_sequence": None},
+            "usage": {"output_tokens": 100},
+        },
+        {"type": "message_stop"},
+    ]
+    for event in events:
+        yield f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
+
+
+def record_parsed_json_sizes(monkeypatch: pytest.MonkeyPatch) -> list[int]:
+    """Patches `jiter.from_json` to record the size of every input it parses."""
+    sizes: list[int] = []
+    from_json = jiter.from_json
+
+    def recording_from_json(json_data: bytes, **kwargs: Any) -> Any:
+        sizes.append(len(json_data))
+        return from_json(json_data, **kwargs)
+
+    monkeypatch.setattr(jiter, "from_json", recording_from_json)
+    return sizes
+
+
+def assert_long_tool_input_parsed(snapshots: list[object], final_input: object, parsed_sizes: list[int]) -> None:
+    # re-parsing the whole buffer on every delta would add up to hundreds of times the input size
+    assert sum(parsed_sizes) < 10 * len(LONG_TOOL_INPUT)
+
+    expected: list[object] = []
+    for i in range(8, len(LONG_TOOL_INPUT), 8):
+        expected.append(jiter.from_json(LONG_TOOL_INPUT[:i].encode(), partial_mode=True))
+    expected.append(json.loads(LONG_TOOL_INPUT))
+
+    assert snapshots == expected
+    assert final_input == json.loads(LONG_TOOL_INPUT)

tests/lib/streaming/test_beta_messages.py

@@ -26,7 +26,13 @@
 from anthropic.types.beta.beta_message_delta_usage import BetaMessageDeltaUsage
 from anthropic.types.beta.beta_raw_message_delta_event import Delta as BetaRawMessageDelta, BetaRawMessageDeltaEvent
 
-from .helpers import get_response, to_async_iter
+from .helpers import (
+    get_response,
+    to_async_iter,
+    record_parsed_json_sizes,
+    get_long_tool_input_response,
+    assert_long_tool_input_parsed,
+)
 
 base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010")
 api_key = "my-anthropic-api-key"
@@ -474,6 +480,22 @@ def test_tool_use(self, respx_mock: MockRouter) -> None:
 
         assert json.loads(respx_mock.calls.last.request.content)["tools"] == [WeatherTool().to_dict()]
 
+    @pytest.mark.respx(base_url=base_url)
+    def test_tool_use_long_string_input(self, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch) -> None:
+        respx_mock.post("/v1/messages").mock(return_value=httpx2.Response(200, content=get_long_tool_input_response()))
+        parsed_sizes = record_parsed_json_sizes(monkeypatch)
+
+        with sync_client.beta.messages.stream(
+            max_tokens=1024,
+            messages=[{"role": "user", "content": "Write the notes file"}],
+            model="claude-sonnet-5-5",
+        ) as stream:
+            snapshots = [event.snapshot for event in stream if event.type == "input_json"]
+            tool_use = stream.get_final_message().content[0]
+
+        assert tool_use.type == "tool_use"
+        assert_long_tool_input_parsed(snapshots, tool_use.input, parsed_sizes)
+
     @pytest.mark.respx(base_url=base_url)
     def test_tool_use_round_trip(self, respx_mock: MockRouter) -> None:
         route = respx_mock.post("/v1/messages").mock(
@@ -838,6 +860,25 @@ async def test_tool_use(self, respx_mock: MockRouter) -> None:
 
         assert json.loads(respx_mock.calls.last.request.content)["tools"] == [WeatherTool().to_dict()]
 
+    @pytest.mark.asyncio
+    @pytest.mark.respx(base_url=base_url)
+    async def test_tool_use_long_string_input(self, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch) -> None:
+        respx_mock.post("/v1/messages").mock(
+            return_value=httpx2.Response(200, content=to_async_iter(get_long_tool_input_response()))
+        )
+        parsed_sizes = record_parsed_json_sizes(monkeypatch)
+
+        async with async_client.beta.messages.stream(
+            max_tokens=1024,
+            messages=[{"role": "user", "content": "Write the notes file"}],
+            model="claude-sonnet-5-5",
+        ) as stream:
+            snapshots = [event.snapshot async for event in stream if event.type == "input_json"]
+            tool_use = (await stream.get_final_message()).content[0]
+
+        assert tool_use.type == "tool_use"
+        assert_long_tool_input_parsed(snapshots, tool_use.input, parsed_sizes)
+
     @pytest.mark.asyncio
     @pytest.mark.respx(base_url=base_url)
     async def test_server_tool_use(self, respx_mock: MockRouter) -> None: