Skip to content

perf(tools): scan only the chunk for a collecting tool call's end tag - #625

Open
spokvulcan wants to merge 1 commit into
ml-explore:mainfrom
spokvulcan:perf/tool-call-end-tag-scan
Open

spokvulcan wants to merge 1 commit into
ml-explore:mainfrom
spokvulcan:perf/tool-call-end-tag-scan

Conversation

@spokvulcan

Copy link
Copy Markdown
Contributor

Proposed changes

Fixes #624.

While collecting a tagged tool call, ToolCallProcessor and the text recovery scanner searched the whole buffered call for its closing tag on every chunk, which is quadratic in the call's length: a 13k-token write call spent seconds in that scan alone. Only the text a chunk appends can complete the closing tag, because an earlier occurrence was already in the buffer when the previous chunk was scanned. Both scanners now check the appended bytes plus the tag's overlap with what preceded them, through one helper on ToolCallFrameScanner, and run the structural frame scan only when that window holds the tag. The chunk that opens the frame still scans the whole buffer, which then holds at most the start tag and that chunk. Five tests in ToolCallProcessorLongCallTests: a 40k-character call streamed one character at a time parses in linear time, with and without declared tools, and the closing tag is found when it is split across chunks, when it arrives in the chunk that completes the start tag, and when another call follows it in the same chunk.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

AI usage

  • I have read this PR description in full and approve it as my own, and it
    accurately describes the code changes.
  • AI usage disclosure: the change, its tests and this description were drafted with Claude Code and reviewed, verified and approved by the author.

`ToolCallProcessor` checked the whole buffered call for its end tag on
every chunk while collecting a tagged call, on the native path and on the
cross-dialect recovery path that declared tools enable, so a call costs
time quadratic in its length: seconds for a call of ten thousand tokens,
paid once on the live stream and again wherever the stream is replayed.

Only the text a chunk appends can complete the end tag — an earlier
occurrence was already in the buffer when the previous chunk was scanned —
so both scanners now check the appended bytes plus the tag's overlap with
what preceded them, and run the full structural frame scan only when that
window holds the tag. The chunk that opens the frame still scans the whole
buffer, which then holds at most the start tag and that chunk.

Tests: a 40k-character call streamed one character at a time parses in
linear time, with and without declared tools; the end tag is found when it
is split across chunks, when it arrives in the chunk that completes the
start tag, and when another call follows it in the same chunk.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] ToolCallProcessor rescans the whole collected call for its end tag on every chunk

2 participants