diff --git a/python/sglang/srt/function_call/kimik2_detector.py b/python/sglang/srt/function_call/kimik2_detector.py
index 2345817de..11909fa6f 100644
--- a/python/sglang/srt/function_call/kimik2_detector.py
+++ b/python/sglang/srt/function_call/kimik2_detector.py
@@ -15,7 +15,6 @@ from sglang.srt.function_call.core_types import (
ToolCallItem,
_GetInfoFunc,
)
-from sglang.srt.function_call.utils import _is_complete_json
logger = logging.getLogger(__name__)
@@ -177,9 +176,15 @@ class KimiK2Detector(BaseFormatDetector):
logger.debug("function_call_tuples: %s", function_call_tuples)
tool_calls = []
+ # ``tool_index`` is the per-response 0-based position of the call
+ # (OpenAI spec); enumerate parsed calls locally and ignore the
+ # model's ``:N`` suffix, which is a conversation-level counter.
+ # ``serving_chat._process_tool_call_id()`` later offsets these by
+ # ``history_tool_calls_cnt`` for multi-turn responses.
+ local_tool_index = 0
for match in function_call_tuples:
function_id, function_args = match
- function_name, function_idx = self._parse_tool_call_id(
+ function_name, _ = self._parse_tool_call_id(
function_id, tools, function_args
)
if function_name is None:
@@ -189,11 +194,12 @@ class KimiK2Detector(BaseFormatDetector):
tool_calls.append(
ToolCallItem(
- tool_index=function_idx,
+ tool_index=local_tool_index,
name=function_name,
parameters=function_args,
)
)
+ local_tool_index += 1
content = text[: text.find(self.bot_token)]
return StreamingParseResult(normal_text=content, calls=tool_calls)
@@ -205,127 +211,207 @@ class KimiK2Detector(BaseFormatDetector):
def parse_streaming_increment(
self, new_text: str, tools: List[Tool]
) -> StreamingParseResult:
- """
- Streaming incremental parsing tool calls for KimiK2 format.
- """
+ """Streaming incremental parsing tool calls for KimiK2 format."""
self._buffer += new_text
- current_text = self._buffer
- # Check if we have a tool call (either the start token or individual tool call)
- has_tool_call = (
- self.bot_token in current_text or self.tool_call_start_token in current_text
- )
-
- if not has_tool_call:
- self._buffer = ""
- normal_text = _strip_special_tokens(new_text)
- return StreamingParseResult(normal_text=normal_text)
+ # Fast path: no tool call in flight and no markers yet -- emit as
+ # normal text, holding back any trailing partial start token.
+ if (
+ self._current_stream_function_name is None
+ and self.bot_token not in self._buffer
+ and self.tool_call_start_token not in self._buffer
+ ):
+ emit, hold = self._split_pending_start(self._buffer)
+ self._buffer = hold
+ return StreamingParseResult(normal_text=_strip_special_tokens(emit))
if not hasattr(self, "_tool_indices"):
self._tool_indices = self._get_tool_indices(tools)
+ normal_text_parts: list[str] = []
calls: list[ToolCallItem] = []
+
try:
- match = self.stream_tool_call_portion_regex.search(current_text)
- if match:
- function_id = match.group("tool_call_id")
- function_args = match.group("function_arguments")
+ while True:
+ buffer = self._buffer
- # Reuse cached name for current tool call to avoid repeated
- # json.loads on partial JSON in _infer_tool_name.
- if self._current_stream_function_name is not None:
- function_name = self._current_stream_function_name
- else:
- function_name, _ = self._parse_tool_call_id(
- function_id, tools, function_args
+ # Locate next <|tool_call_begin|>, draining any prefix as text.
+ begin_idx = self._locate_tool_call_start(buffer, normal_text_parts)
+ if begin_idx is None:
+ break
+ buffer = self._buffer
+
+ # If another <|tool_call_begin|> appears before the header
+ # closes with <|tool_call_argument_begin|>, the section is
+ # malformed -- discard and restart at the orphan.
+ arg_begin_idx = buffer.find(self.tool_call_argument_begin_token)
+ next_begin = buffer.find(
+ self.tool_call_start_token, len(self.tool_call_start_token)
+ )
+ if next_begin != -1 and (
+ arg_begin_idx == -1 or next_begin < arg_begin_idx
+ ):
+ logger.warning(
+ "Kimi-K2 tool_call_begin without preceding tool_call_end; "
+ "discarding incomplete section."
)
- if function_name is None:
- return StreamingParseResult(normal_text="", calls=calls)
+ self._buffer = buffer[next_begin:]
+ self._reset_inflight_call_state()
+ continue
- # Initialize state if this is the first tool call
- if self.current_tool_id == -1:
- self.current_tool_id = 0
- self.prev_tool_call_arr = []
- self.streamed_args_for_tool = [""]
+ if arg_begin_idx == -1:
+ # Header not fully arrived yet.
+ break
- # Ensure we have enough entries in our tracking arrays
- while len(self.prev_tool_call_arr) <= self.current_tool_id:
- self.prev_tool_call_arr.append({})
- while len(self.streamed_args_for_tool) <= self.current_tool_id:
- self.streamed_args_for_tool.append("")
+ id_start = len(self.tool_call_start_token)
+ function_id = buffer[id_start:arg_begin_idx].strip()
+ args_start = arg_begin_idx + len(self.tool_call_argument_begin_token)
+ end_idx = buffer.find(self.tool_call_end_token)
- if not self.current_tool_name_sent:
+ # Resolve function name (cached across chunks within a section).
+ name_just_resolved = False
+ if self._current_stream_function_name is None:
+ args_for_inference = (
+ buffer[args_start:end_idx]
+ if end_idx != -1
+ else buffer[args_start:]
+ )
+ resolved = self._resolve_function_name(
+ function_id, tools, args_for_inference
+ )
+ if resolved is None:
+ if end_idx == -1:
+ # Wait for the end marker before deciding.
+ break
+ logger.warning(
+ "Kimi-K2 unrecognized tool_call_id %r; skipping section.",
+ function_id,
+ )
+ self._buffer = buffer[end_idx + len(self.tool_call_end_token) :]
+ self._reset_inflight_call_state()
+ continue
+ name = resolved
+ self._current_stream_function_name = name
+ name_just_resolved = True
+
+ # ``tool_index`` is the per-response 0-based position
+ # (OpenAI streaming spec); ignore the model's ``:N`` suffix
+ # which is a conversation-level counter.
+ if self.current_tool_id == -1:
+ self.current_tool_id = 0
+ self.prev_tool_call_arr = []
+ self.streamed_args_for_tool = [""]
+ while len(self.prev_tool_call_arr) <= self.current_tool_id:
+ self.prev_tool_call_arr.append({})
+ while len(self.streamed_args_for_tool) <= self.current_tool_id:
+ self.streamed_args_for_tool.append("")
+ self.prev_tool_call_arr[self.current_tool_id] = {
+ "name": name,
+ "arguments": {},
+ }
+ self.current_tool_name_sent = True
+
+ # Stream newly-arrived args, combining the first event with
+ # the freshly-resolved name.
+ if end_idx != -1:
+ args_full = buffer[args_start:end_idx]
+ else:
+ args_full = buffer[args_start:]
+ argument_diff = args_full[len(self._last_arguments) :]
+ if argument_diff or name_just_resolved:
calls.append(
ToolCallItem(
tool_index=self.current_tool_id,
- name=function_name,
- parameters="",
+ name=(
+ self._current_stream_function_name
+ if name_just_resolved
+ else None
+ ),
+ parameters=argument_diff,
)
)
- self.current_tool_name_sent = True
- self._current_stream_function_name = function_name
- self.prev_tool_call_arr[self.current_tool_id] = {
- "name": function_name,
- "arguments": {},
- }
- else:
- argument_diff = (
- function_args[len(self._last_arguments) :]
- if function_args.startswith(self._last_arguments)
- else function_args
- )
-
- parsed_args_diff = argument_diff.split(self.tool_call_end_token, 1)[
- 0
- ]
-
- if parsed_args_diff:
- calls.append(
- ToolCallItem(
- tool_index=self.current_tool_id,
- name=None,
- parameters=parsed_args_diff,
- )
- )
- self._last_arguments += parsed_args_diff
+ if argument_diff:
+ self._last_arguments += argument_diff
self.streamed_args_for_tool[
self.current_tool_id
- ] += parsed_args_diff
+ ] += argument_diff
- parsed_args = function_args.split(self.tool_call_end_token, 1)[0]
- if _is_complete_json(parsed_args):
- try:
- parsed_args = json.loads(parsed_args)
- self.prev_tool_call_arr[self.current_tool_id][
- "arguments"
- ] = parsed_args
- except json.JSONDecodeError:
- pass
+ if end_idx == -1:
+ # Args still streaming.
+ break
- # Find the end of the current tool call and remove only that part from buffer
- tool_call_end_pattern = (
- r"<\|tool_call_begin\|>.*?<\|tool_call_end\|>"
- )
- end_match = re.search(
- tool_call_end_pattern, current_text, re.DOTALL
- )
- if end_match:
- self._buffer = current_text[end_match.end() :]
- else:
- self._buffer = ""
+ # Section finalized -- advance buffer and prepare next call.
+ self._buffer = buffer[end_idx + len(self.tool_call_end_token) :]
+ self.current_tool_id += 1
+ self._reset_inflight_call_state()
- result = StreamingParseResult(normal_text="", calls=calls)
- self.current_tool_id += 1
- self._last_arguments = ""
- self.current_tool_name_sent = False
- self._current_stream_function_name = None
- return result
-
- return StreamingParseResult(normal_text="", calls=calls)
+ return StreamingParseResult(
+ normal_text="".join(normal_text_parts), calls=calls
+ )
except Exception as e:
logger.error("Error in parse_streaming_increment: %s", e, exc_info=True)
- return StreamingParseResult(normal_text=_strip_special_tokens(current_text))
+ # Drop the buffer to avoid leaking raw special tokens.
+ self._buffer = ""
+ self._reset_inflight_call_state()
+ return StreamingParseResult(
+ normal_text="".join(normal_text_parts), calls=calls
+ )
+
+ def _reset_inflight_call_state(self) -> None:
+ """Reset per-section streaming state after finalize/discard."""
+ self._last_arguments = ""
+ self.current_tool_name_sent = False
+ self._current_stream_function_name = None
+
+ def _locate_tool_call_start(
+ self, buffer: str, normal_text_parts: list
+ ) -> int | None:
+ """Find the next <|tool_call_begin|>; drain any prefix as normal text.
+
+ Returns 0 on success, or ``None`` when no start token is present yet.
+ """
+ begin_idx = buffer.find(self.tool_call_start_token)
+ if begin_idx == -1:
+ emit, hold = self._split_pending_start(buffer)
+ if emit:
+ normal_text_parts.append(_strip_special_tokens(emit))
+ self._buffer = hold
+ return None
+
+ if begin_idx > 0:
+ normal_text_parts.append(_strip_special_tokens(buffer[:begin_idx]))
+ self._buffer = buffer[begin_idx:]
+ return 0
+
+ def _split_pending_start(self, text: str) -> tuple[str, str]:
+ """Hold back a trailing fragment that could be the start of
+ <|tool_calls_section_begin|> or <|tool_call_begin|>. Everything
+ before it is safe to emit as normal text.
+ """
+ candidates = (self.bot_token, self.tool_call_start_token)
+ max_tail = max(len(t) for t in candidates) - 1
+ for n in range(min(len(text), max_tail), 1, -1):
+ tail = text[-n:]
+ if any(t.startswith(tail) for t in candidates):
+ return text[:-n], tail
+ return text, ""
+
+ def _resolve_function_name(
+ self, function_id: str, tools: List[Tool], function_args: str
+ ) -> Optional[str]:
+ """Map a Kimi-K2 tool_call_id to a tool name, or ``None`` if unknown."""
+ if not function_id:
+ return self._infer_tool_name(tools, function_args)
+
+ m = self.tool_call_id_regex.match(function_id)
+ if m:
+ return m.group("name")
+
+ if self.tool_call_id_counter_regex.match(function_id):
+ return self._infer_tool_name(tools, function_args)
+
+ return None
def structure_info(self) -> _GetInfoFunc:
"""Return function that creates StructureInfo for guided generation."""
diff --git a/test/registered/function_call/test_kimik2_detector.py b/test/registered/function_call/test_kimik2_detector.py
index ee1cccf16..736ed1624 100644
--- a/test/registered/function_call/test_kimik2_detector.py
+++ b/test/registered/function_call/test_kimik2_detector.py
@@ -111,6 +111,32 @@ class TestKimiK2DetectorBasic(unittest.TestCase):
self.assertEqual(result.calls[1].name, "get_weather")
self.assertEqual(result.calls[1].parameters, '{"city": "Tokyo"}')
+ def test_non_streaming_tool_index_is_local(self):
+ """tool_index is the per-response 0-based position, not the model's :N suffix.
+
+ The model may emit conversation-level ``:N`` counters (e.g. ``:5``, ``:6``)
+ in a multi-turn conversation. The non-streaming parser must enumerate
+ parsed calls locally (0, 1, ...) so that
+ ``serving_chat._process_tool_call_id()`` can offset them by
+ ``history_tool_calls_cnt`` without double-counting.
+ """
+ text = (
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.ReadFile:5"
+ '<|tool_call_argument_begin|>{"path": "/a.py"}'
+ "<|tool_call_end|>"
+ "<|tool_call_begin|>functions.get_weather:6"
+ '<|tool_call_argument_begin|>{"city": "Tokyo"}'
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ )
+ result = self.detector.detect_and_parse(text, self.tools)
+ self.assertEqual(len(result.calls), 2)
+ self.assertEqual(result.calls[0].tool_index, 0)
+ self.assertEqual(result.calls[0].name, "ReadFile")
+ self.assertEqual(result.calls[1].tool_index, 1)
+ self.assertEqual(result.calls[1].name, "get_weather")
+
def test_normal_text_before_tool_call(self):
"""Normal text before tool call markers is preserved."""
text = (
@@ -183,6 +209,15 @@ class TestKimiK2DetectorHyphenatedNames(unittest.TestCase):
class TestKimiK2DetectorStreaming(unittest.TestCase):
"""Streaming incremental parsing tests for KimiK2Detector."""
+ def test_streaming_trailing_literal_left_angle_is_not_dropped(self):
+ """A final literal '<' must remain in normal_text instead of being buffered away."""
+ detector = KimiK2FuncDetector()
+
+ result = detector.parse_streaming_increment("normal text <", [])
+
+ self.assertEqual(result.normal_text, "normal text <")
+ self.assertEqual(detector._buffer, "")
+
def setUp(self):
self.tools = [
_make_tool("ReadFile"),
@@ -627,6 +662,409 @@ class TestKimiK2EndToEnd(unittest.TestCase):
self.assertEqual(len(name_calls), 1)
self.assertEqual(name_calls[0].name, "get_weather")
+ def test_e2e_normal_think_close_then_content_overlap_tool_call(self):
+ """Trailing content between ```` and the tool-call markers
+ (e.g. ``"This is a content:"``) must be surfaced as ``normal_text``
+ by the tool-call parser and not stripped — this is the exact bug the
+ PR fixes.
+ """
+ reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
+ tc_det = KimiK2FuncDetector()
+
+ chunks = [
+ "",
+ "Thinking about it...",
+ "This is a ",
+ "content:<|tool_calls_section_begin|>",
+ "<|tool_call_begin|>functions.get_weather:0",
+ '<|tool_call_argument_begin|>{"city": "London"}',
+ "<|tool_call_end|>",
+ "<|tool_calls_section_end|>",
+ ]
+
+ all_reasoning = ""
+ all_content = ""
+
+ toolcall_chunks = []
+ for chunk in chunks:
+ r = reasoning_det.parse_streaming_increment(chunk)
+ all_reasoning += r.reasoning_text
+ if r.normal_text:
+ toolcall_chunks.append(r.normal_text)
+
+ tool_calls, all_content = _collect_streaming_tool_calls(
+ tc_det, toolcall_chunks, self.tools
+ )
+
+ self.assertEqual("Thinking about it...", all_reasoning)
+ self.assertEqual("This is a content:", all_content)
+ self.assertEqual(len(tool_calls), 1)
+ first_call = tool_calls.pop()
+ self.assertEqual(first_call["name"], "get_weather")
+ self.assertEqual(first_call["parameters"], '{"city": "London"}')
+
+ def test_e2e_normal_think_close_then_content_overlap_tool_call_multi_token(self):
+ """Speculative decoding: a single chunk may contain normal text followed
+ by tool-call markers and even the tool_call_begin/id. The normal-text
+ prefix must still be emitted (not stripped) by the tool-call parser.
+ """
+ reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
+ tc_det = KimiK2FuncDetector()
+
+ chunks = [
+ "",
+ "Thinking about it...",
+ "This is a ",
+ "content:<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0"
+ '<|tool_call_argument_begin|>{"city": "London"}<|tool_call_end|><|tool_calls_section_end|>',
+ ]
+
+ all_reasoning = ""
+ all_content = ""
+ toolcall_chunks = []
+
+ for chunk in chunks:
+ r = reasoning_det.parse_streaming_increment(chunk)
+ all_reasoning += r.reasoning_text
+ if r.normal_text:
+ toolcall_chunks.append(r.normal_text)
+
+ tool_calls, all_content = _collect_streaming_tool_calls(
+ tc_det, toolcall_chunks, self.tools
+ )
+
+ self.assertEqual("Thinking about it...", all_reasoning)
+ self.assertEqual("This is a content:", all_content)
+ self.assertEqual(len(tool_calls), 1)
+ first_call = tool_calls.pop()
+ self.assertEqual(first_call["name"], "get_weather")
+ self.assertEqual(first_call["parameters"], '{"city": "London"}')
+
+ def test_e2e_normal_think_close_then_content_overlap_tool_call_multi_token_multi_calls(
+ self,
+ ):
+ """Speculative decoding: a single chunk may contain normal text followed
+ by tool-call markers and even the tool_call_begin/id. Additionally, this
+ single chunk packs two complete tool-call sections back-to-back, which
+ exercises the ``while True:`` drain loop introduced by this PR — both
+ calls must be emitted from one ``parse_streaming_increment`` invocation.
+ """
+ reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
+ tc_det = KimiK2FuncDetector()
+
+ chunks = [
+ "",
+ "Thinking about it...",
+ "This is a ",
+ 'content:<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"city": "London"}<|tool_call_end|><|tool_calls_section_end|><|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:1<|tool_call_argument_begin|>'
+ '{"city": "Delhi"}<|tool_call_end|><|tool_calls_section_end|>',
+ ]
+
+ all_reasoning = ""
+ all_content = ""
+ toolcall_chunks = []
+
+ for chunk in chunks:
+ r = reasoning_det.parse_streaming_increment(chunk)
+ all_reasoning += r.reasoning_text
+ if r.normal_text:
+ toolcall_chunks.append(r.normal_text)
+
+ tool_calls, all_content = _collect_streaming_tool_calls(
+ tc_det, toolcall_chunks, self.tools
+ )
+
+ self.assertEqual("Thinking about it...", all_reasoning)
+ self.assertEqual("This is a content:", all_content)
+ self.assertEqual(len(tool_calls), 2)
+ first_call = tool_calls.pop(0)
+ self.assertEqual(first_call["name"], "get_weather")
+ self.assertEqual(first_call["parameters"], '{"city": "London"}')
+ second_call = tool_calls.pop(0)
+ self.assertEqual(second_call["name"], "get_weather")
+ self.assertEqual(second_call["parameters"], '{"city": "Delhi"}')
+
+ def test_e2e_chunk_split_invariance(self):
+ """The detector must produce identical results across a few realistic
+ chunking variants. Special tokens (e.g. ``<|tool_calls_section_begin|>``)
+ are atomic and never split, so cuts only fall on token boundaries or
+ inside JSON args.
+ """
+ prefix = "Thinking about it...This is a content:"
+ call1 = (
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:0"
+ '<|tool_call_argument_begin|>{"city": "London"}'
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ )
+ call2 = (
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:1"
+ '<|tool_call_argument_begin|>{"city": "Delhi"}'
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ )
+
+ expected_reasoning = "Thinking about it..."
+ expected_content = "This is a content:"
+ expected_calls = [
+ {"name": "get_weather", "parameters": '{"city": "London"}'},
+ {"name": "get_weather", "parameters": '{"city": "Delhi"}'},
+ ]
+
+ variants = {
+ # One complete tool call per chunk.
+ "first_complete_then_second_complete": [prefix + call1, call2],
+ # Both tool calls arrive in a single chunk.
+ "both_in_one_chunk": [prefix + call1 + call2],
+ # First call complete + second call partial (cut inside JSON args),
+ # then the rest of the second call.
+ "first_complete_second_partial_then_rest": [
+ prefix
+ + call1
+ + "<|tool_calls_section_begin|>"
+ + "<|tool_call_begin|>functions.get_weather:1"
+ + '<|tool_call_argument_begin|>{"city": "De',
+ 'lhi"}<|tool_call_end|><|tool_calls_section_end|>',
+ ],
+ # First call partial (cut inside JSON args), then rest of first +
+ # full second call.
+ "first_partial_then_first_complete_second_complete": [
+ prefix
+ + "<|tool_calls_section_begin|>"
+ + "<|tool_call_begin|>functions.get_weather:0"
+ + '<|tool_call_argument_begin|>{"city": "Lon',
+ 'don"}<|tool_call_end|><|tool_calls_section_end|>' + call2,
+ ],
+ }
+
+ for name, chunks in variants.items():
+ with self.subTest(variant=name):
+ reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
+ tc_det = KimiK2FuncDetector()
+ all_reasoning = ""
+ toolcall_chunks = []
+ for chunk in chunks:
+ r = reasoning_det.parse_streaming_increment(chunk)
+ all_reasoning += r.reasoning_text
+ if r.normal_text:
+ toolcall_chunks.append(r.normal_text)
+ tool_calls, all_content = _collect_streaming_tool_calls(
+ tc_det, toolcall_chunks, self.tools
+ )
+ self.assertEqual(all_reasoning, expected_reasoning)
+ self.assertEqual(all_content, expected_content)
+ self.assertEqual(tool_calls, expected_calls)
+
+ def test_e2e_normal_text_between_two_tool_calls(self):
+ """Normal text appearing BETWEEN two tool-call sections must be
+ surfaced as ``normal_text``.
+ """
+ tc_det = KimiK2FuncDetector()
+ chunks = [
+ "Prefix text:"
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:0"
+ '<|tool_call_argument_begin|>{"city": "London"}'
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ " Now calling next: "
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:1"
+ "<|tool_call_argument_begin|>"
+ '{"city":'
+ ' "Delhi"}'
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ ]
+
+ tool_calls, all_content = _collect_streaming_tool_calls(
+ tc_det, chunks, self.tools
+ )
+
+ self.assertIn("Prefix text:", all_content)
+ self.assertIn("Now calling next:", all_content)
+ self.assertEqual(len(tool_calls), 2)
+ self.assertEqual(tool_calls[0]["name"], "get_weather")
+ self.assertEqual(tool_calls[0]["parameters"], '{"city": "London"}')
+ self.assertEqual(tool_calls[1]["name"], "get_weather")
+ self.assertEqual(tool_calls[1]["parameters"], '{"city": "Delhi"}')
+
+ def test_e2e_unparsable_tool_id_does_not_wedge_stream(self):
+ """A tool_call header with an unparsable ID must not wedge the
+ streaming parser.
+ """
+ tc_det = KimiK2FuncDetector()
+
+ # ``weird@id`` matches the broad ``[^\\s<|]+`` capture in
+ # ``stream_tool_call_portion_regex`` but fails both the standard
+ # ``name:idx`` form and the bare-counter form, so
+ # ``_parse_tool_call_id`` returns ``(None, 0)``.
+ chunks = [
+ "normal text before",
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>weird@id"
+ '<|tool_call_argument_begin|>{"city"'
+ ': "London"}'
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>",
+ # A valid follow-up call: must still be parsed.
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:0"
+ '<|tool_call_argument_begin|>{"city": "Delhi"}'
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>",
+ ]
+
+ tool_calls, all_content = _collect_streaming_tool_calls(
+ tc_det, chunks, self.tools
+ )
+
+ # The bad call may be skipped/logged, but the follow-up MUST be
+ # emitted. The stream must not be wedged on the bad header.
+ valid = [c for c in tool_calls if c.get("name") == "get_weather"]
+ self.assertEqual(len(valid), 1)
+ self.assertEqual(valid[0]["parameters"], '{"city": "Delhi"}')
+ self.assertIn("normal text before", all_content)
+
+ def test_e2e_malformed_json_args_passes_through_to_client(self):
+ """A tool call with malformed JSON args (e.g. unclosed brace,
+ spelling mistake) is the **client's** problem to
+ validate/repair — the parser's job is to locate boundaries and
+ hand back the raw argument string. This mirrors the
+ ``detect_and_parse`` (non-streaming) contract.
+
+ Required behavior:
+
+ 1. The malformed call is emitted unchanged (raw bytes
+ preserved, name + index intact) so the client can decide
+ how to handle it (reject, repair, replay-prompt, etc).
+ 2. The stream is not wedged — the trailing valid call must
+ still parse.
+ 3. ``current_tool_id`` advances normally — the trailing valid
+ call sits at index 2, not 1 or 3.
+
+ Asserts the SAME outcome under three chunk layouts:
+
+ * **single-chunk / MTP path** — all three sections in one
+ forward step (mimics speculative / multi-token-prediction).
+ * **per-call split path** — one section per chunk.
+ * **bad section split mid-payload** — the malformed section
+ itself is fragmented across three chunks (header + partial
+ args; more args; end-token + trailing valid call). Verifies
+ the atomic-section buffer correctly defers emission until
+ ``<|tool_call_end|>`` arrives.
+ """
+ good_args_0 = '{"city": "London"}'
+ # JSON keyword (the model misspelled ``false``).
+ bad_args = '{"city": "Bad", "valid": fasle'
+ good_args_1 = '{"city": "Delhi"}'
+
+ good_section_0 = (
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:0"
+ f"<|tool_call_argument_begin|>{good_args_0}"
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ )
+ bad_section = (
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:1"
+ f"<|tool_call_argument_begin|>{bad_args}"
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ )
+ good_section_1 = (
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:2"
+ f"<|tool_call_argument_begin|>{good_args_1}"
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ )
+
+ layouts = {
+ "mtp_single_chunk": [good_section_0 + bad_section + good_section_1],
+ "per_call_chunks": [good_section_0, bad_section, good_section_1],
+ "bad_section_split_mid_payload": [
+ good_section_0,
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:1"
+ f'<|tool_call_argument_begin|>{{"city":',
+ ' "Bad", "valid": fasle',
+ "<|tool_call_end|>" "<|tool_calls_section_end|>" + good_section_1,
+ ],
+ }
+
+ for layout_name, chunks in layouts.items():
+ with self.subTest(layout=layout_name):
+ tc_det = KimiK2FuncDetector()
+ tool_calls, _ = _collect_streaming_tool_calls(
+ tc_det, chunks, self.tools
+ )
+
+ # Three contiguous tool calls (good, bad-passthrough, good).
+ self.assertEqual(
+ len(tool_calls),
+ 3,
+ f"[{layout_name}] expected 3 calls, got: {tool_calls!r}",
+ )
+ self.assertEqual(tool_calls[0]["name"], "get_weather")
+ self.assertEqual(tool_calls[1]["name"], "get_weather")
+ self.assertEqual(tool_calls[2]["name"], "get_weather")
+ self.assertEqual(tool_calls[0]["parameters"], good_args_0)
+ # Bad payload preserved byte-for-byte for the client.
+ self.assertEqual(tool_calls[1]["parameters"], bad_args)
+ self.assertEqual(tool_calls[2]["parameters"], good_args_1)
+
+ def test_e2e_exception_mid_drain_preserves_accumulated_calls(self):
+ """An exception raised mid-drain must not discard tool calls already
+ finalized in the same ``parse_streaming_increment`` invocation.
+ """
+ import unittest.mock as mock
+
+ tc_det = KimiK2FuncDetector()
+
+ chunks = [
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:0"
+ '<|tool_call_argument_begin|>{"city": "London"}'
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:1"
+ '<|tool_call_argument_begin|>{"city": "Delhi"}'
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ ]
+
+ # Force an exception on the SECOND resolution step by making
+ # ``_resolve_function_name`` raise the second time it is invoked.
+ call_count = {"n": 0}
+ real_resolve = tc_det._resolve_function_name
+
+ def flaky_resolve(function_id, tools, function_args=None):
+ call_count["n"] += 1
+ if call_count["n"] >= 2:
+ raise RuntimeError("forced mid-drain failure")
+ return real_resolve(function_id, tools, function_args)
+
+ with mock.patch.object(
+ tc_det, "_resolve_function_name", side_effect=flaky_resolve
+ ):
+ tool_calls, all_content = _collect_streaming_tool_calls(
+ tc_det, chunks, self.tools
+ )
+
+ # First call must survive the mid-drain exception.
+ named = [c for c in tool_calls if c.get("name")]
+ self.assertGreaterEqual(len(named), 1)
+ self.assertEqual(named[0]["name"], "get_weather")
+ self.assertEqual(named[0]["parameters"], '{"city": "London"}')
+ # Already-finalized call payload must NOT leak into normal_text.
+ self.assertNotIn('{"city": "London"}', all_content)
+ self.assertNotIn("<|tool_call_begin|>", all_content)
+
def test_e2e_multiple_tool_calls_without_think_close(self):
"""Multiple tool calls inside without ."""
reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
@@ -664,6 +1102,155 @@ class TestKimiK2EndToEnd(unittest.TestCase):
self.assertEqual(name_calls[1].name, "ReadFile")
+# ============================================================
+# Part 2b: OpenAI streaming-spec compliance for ``tool_index``
+# ============================================================
+
+
+class TestKimiK2DetectorOpenAIIndexCompliance(unittest.TestCase):
+ """The detector must emit ``tool_index`` as a dense, 0-based position
+ within the *current response* (per the OpenAI streaming spec), regardless
+ of the value of the model's conversation-level ``:N`` counter in the
+ tool_call header. The serving layer is responsible for adding any
+ history offset back when synthesizing the public ``id`` field.
+
+ These tests pin the detector contract so multi-turn conversations
+ (where the model continues an auto-incrementing counter across turns)
+ can never produce sparse / non-zero-based ``index`` values in the
+ streamed delta.
+ """
+
+ def setUp(self):
+ self.tools = [
+ _make_tool(
+ "get_weather",
+ {
+ "type": "object",
+ "properties": {"city": {"type": "string"}},
+ "required": ["city"],
+ },
+ ),
+ ]
+
+ def _run(self, chunks):
+ det = KimiK2FuncDetector()
+ events = []
+ for chunk in chunks:
+ r = det.parse_streaming_increment(chunk, self.tools)
+ events.extend(r.calls)
+ return events
+
+ def test_single_call_with_nonzero_model_counter_starts_at_index_0(self):
+ """Model continues a conversation-level counter (``:5``) across turns.
+ The detector must still emit ``tool_index=0`` for the first call in
+ this response — the model's ``:N`` suffix MUST NOT leak into ``index``.
+ """
+ chunks = [
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:5"
+ '<|tool_call_argument_begin|>{"city":',
+ ' "Paris"}<|tool_call_end|><|tool_calls_section_end|>',
+ ]
+ events = self._run(chunks)
+
+ self.assertGreaterEqual(len(events), 1)
+ # Every emitted delta (name event + arg deltas) must use index 0.
+ for ev in events:
+ self.assertEqual(
+ ev.tool_index,
+ 0,
+ f"tool_index must be 0-based per response, got {ev.tool_index}",
+ )
+ first = events[0]
+ self.assertEqual(first.name, "get_weather")
+
+ def test_multi_call_response_uses_dense_0_based_indices(self):
+ """Two calls in one response, model emits ``:7`` then ``:8``. The
+ detector must emit ``tool_index=0`` then ``tool_index=1`` (dense,
+ 0-based), independent of the model's counter.
+ """
+ chunks = [
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:7"
+ '<|tool_call_argument_begin|>{"city": "London"}'
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:8"
+ '<|tool_call_argument_begin|>{"city": "Berlin"}'
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ ]
+ events = self._run(chunks)
+
+ name_events = [e for e in events if e.name]
+ self.assertEqual(len(name_events), 2)
+ self.assertEqual(name_events[0].tool_index, 0)
+ self.assertEqual(name_events[1].tool_index, 1)
+
+ # Every delta for call N must carry tool_index == N.
+ for ev in events:
+ self.assertIn(ev.tool_index, (0, 1))
+
+ def test_continuation_chunks_keep_same_tool_index(self):
+ """Per OpenAI spec, all argument-delta chunks for a given call
+ must share the same ``index``. Split the args across multiple
+ chunks and verify ``tool_index`` stays at 0 throughout.
+ """
+ chunks = [
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:10"
+ '<|tool_call_argument_begin|>{"city":',
+ ' "Pa',
+ 'ris"}',
+ "<|tool_call_end|><|tool_calls_section_end|>",
+ ]
+ events = self._run(chunks)
+
+ self.assertGreater(len(events), 1, "expected name event + arg deltas")
+ for ev in events:
+ self.assertEqual(ev.tool_index, 0)
+
+ # Reassembled args must round-trip.
+ joined = "".join(ev.parameters or "" for ev in events)
+ self.assertEqual(joined, '{"city": "Paris"}')
+
+ def test_three_calls_with_nonzero_model_counter_indices_are_dense(self):
+ """Multi-turn worst case: model continues at ``:10`` and emits three
+ calls in this response. Indices must be 0, 1, 2 — not 10, 11, 12.
+ """
+ sections = []
+ for offset, city in enumerate(("Paris", "Berlin", "Madrid")):
+ sections.append(
+ "<|tool_calls_section_begin|>"
+ f"<|tool_call_begin|>functions.get_weather:{10 + offset}"
+ f'<|tool_call_argument_begin|>{{"city": "{city}"}}'
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ )
+ events = self._run(["".join(sections)])
+
+ name_events = [e for e in events if e.name]
+ self.assertEqual([e.tool_index for e in name_events], [0, 1, 2])
+
+ def test_bare_counter_id_also_starts_at_index_0(self):
+ """Same invariant for the bare-counter ID form (model omits function
+ name and emits just a numeric counter, e.g. ``:5``).
+ """
+ chunks = [
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>5"
+ '<|tool_call_argument_begin|>{"city": "Paris"}'
+ "<|tool_call_end|>"
+ "<|tool_calls_section_end|>"
+ ]
+ events = self._run(chunks)
+
+ self.assertGreaterEqual(len(events), 1)
+ for ev in events:
+ self.assertEqual(ev.tool_index, 0)
+
+
# ============================================================
# Part 3: Bare-counter tool call ID parsing
# ============================================================