Fix dropped tool calls when a stream delta carries several (#31860)

2026-07-21 18:07:43 +08:00
parent f369a820d4
commit 2cf2e7377f
2 changed files with 104 additions and 10 deletions
@@ -149,6 +149,68 @@ class TestInklingDetector(unittest.TestCase):
        self.assertEqual(json.loads(args_by_index[0]), {"city": "SF"})
        self.assertEqual(json.loads(args_by_index[1]), {"city": "NY"})

+    def test_streaming_two_complete_tool_calls_in_one_delta_both_emit(self):
+        """Bug regression: parse_streaming_increment parsed one call per delta
+        and re-buffered the rest, relying on the NEXT delta to drain it. Two
+        complete calls arriving in a single (e.g. final) delta left the second
+        stranded in the buffer with no stream-end flush, so only the first was
+        emitted."""
+        detector = InklingDetector()
+        source = (
+            "<|message_model|>weather<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"SF"}}<|end_message|>'
+            "<|message_model|>weather<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"NY"}}<|end_message|>'
+        )
+        args_by_index: dict = {}
+        for call in detector.parse_streaming_increment(source, self.tools).calls:
+            args_by_index[call.tool_index] = (
+                args_by_index.get(call.tool_index, "") + call.parameters
+            )
+        self.assertEqual(sorted(args_by_index), [0, 1])
+        self.assertEqual(json.loads(args_by_index[0]), {"city": "SF"})
+        self.assertEqual(json.loads(args_by_index[1]), {"city": "NY"})
+
+    def test_streaming_text_then_tool_call_in_one_delta_emits_both(self):
+        """Bug regression: a delta carrying visible text followed by a complete
+        tool call emitted only the text and stranded the call in the buffer
+        (the drain loop stopped after the leading-text run), so a final such
+        delta dropped the call. The drain must continue past leading text."""
+        detector = InklingDetector()
+        source = (
+            "Sure, let me check.<|message_model|>weather<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"SF"}}<|end_message|>'
+        )
+        result = detector.parse_streaming_increment(source, self.tools)
+        self.assertIn("Sure, let me check.", result.normal_text)
+        names = [c.name for c in result.calls if c.name]
+        self.assertEqual(names, ["weather"])
+        args = "".join(c.parameters for c in result.calls)
+        self.assertEqual(json.loads(args), {"city": "SF"})
+
+    def test_streaming_rejected_middle_call_keeps_later_valid_call(self):
+        """Bug regression: a rejected call (header/name mismatch) cleared the
+        whole buffer, discarding a later valid call that arrived in the same
+        delta. Only the rejected call's span may be dropped; the drain must
+        continue so the trailing valid call still streams."""
+        detector = InklingDetector()
+        source = (
+            "<|message_model|>weather<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"SF"}}<|end_message|>'
+            "<|message_model|>other<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"XX"}}<|end_message|>'
+            "<|message_model|>weather<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"NY"}}<|end_message|>'
+        )
+        args_by_index: dict = {}
+        for call in detector.parse_streaming_increment(source, self.tools).calls:
+            args_by_index[call.tool_index] = (
+                args_by_index.get(call.tool_index, "") + call.parameters
+            )
+        self.assertEqual(sorted(args_by_index), [0, 1])
+        self.assertEqual(json.loads(args_by_index[0]), {"city": "SF"})
+        self.assertEqual(json.loads(args_by_index[1]), {"city": "NY"})
+
    def test_streaming_rejection_does_not_collide_tool_indices(self):
        """Bug regression: a rejected mid-stream call reset current_tool_id to
        -1, so the NEXT valid call re-announced as tool_index 0 — colliding