Fix json_tools: match "<tool_call" prefix to survive ">" token merge

mlx_lm/tool_parsers/json_tools.py

@@ -2,10 +2,46 @@
 
 import json
 
-tool_call_start = "<tool_call>"
+# The start marker intentionally omits the closing ">". Tool-call markers are matched as
+# exact token-id sequences, and "<tool_call>" encodes to a run ending in a standalone ">"
+# token. Many tokenizers merge that ">" with the following byte (e.g. "<tool_call>\n" -> a
+# single ">\n" token), so the full-marker token run never appears in the generated stream
+# and the call is never captured. Matching the stable "<tool_call" prefix avoids this;
+# parse_tool_call() below tolerates the leftover ">" before the JSON.
+tool_call_start = "<tool_call"
 
 tool_call_end = "</tool_call>"
 
 
 def parse_tool_call(text, tools=None):
-    return json.loads(text.strip())
+    """Extract the first brace-balanced JSON object from the captured tool segment.
+
+    Robust to a leading ">"/newline left by the prefix start marker and to a trailing
+    "</tool_call>" if the end marker likewise merges. Output is identical to
+    json.loads(text.strip()) for a clean JSON-only segment.
+    """
+    start = text.find("{")
+    if start == -1:
+        raise ValueError("no JSON object in tool call segment")
+    depth = 0
+    in_str = False
+    esc = False
+    for i in range(start, len(text)):
+        c = text[i]
+        if in_str:
+            if esc:
+                esc = False
+            elif c == "\\":
+                esc = True
+            elif c == '"':
+                in_str = False
+            continue
+        if c == '"':
+            in_str = True
+        elif c == "{":
+            depth += 1
+        elif c == "}":
+            depth -= 1
+            if depth == 0:
+                return json.loads(text[start : i + 1])
+    raise ValueError("unbalanced JSON in tool call segment")

tests/test_tool_parsing.py

@@ -207,6 +207,28 @@ def test_gemma4(self):
             {"settings": {"enabled": True, "name": "test"}},
         )
 
+    def test_json_tools_marker_merge(self):
+        # The start marker is "<tool_call" (without the closing ">", which many tokenizers
+        # merge with the next byte). The server therefore captures the tool segment starting
+        # at the leftover ">", and the segment may also include a trailing "</tool_call>" if
+        # that marker's tokens merge too. parse_tool_call must extract the JSON regardless.
+        expected = {"name": "get_weather", "arguments": {"city": "Paris"}}
+        variants = [
+            '{"name": "get_weather", "arguments": {"city": "Paris"}}',  # clean (backward-compat)
+            '>\n{"name": "get_weather", "arguments": {"city": "Paris"}}\n',  # leftover ">"
+            '>\n{"name": "get_weather", "arguments": {"city": "Paris"}}\n</tool_call>',  # + end marker
+        ]
+        for text in variants:
+            with self.subTest(text=text):
+                self.assertEqual(json_tools.parse_tool_call(text, None), expected)
+
+        # nested braces and a "}" inside a string value must not truncate extraction
+        text = '{"name": "f", "arguments": {"expr": "a{b}c", "n": {"k": 2}}}'
+        self.assertEqual(
+            json_tools.parse_tool_call(text, None),
+            {"name": "f", "arguments": {"expr": "a{b}c", "n": {"k": 2}}},
+        )
+
         # Array of strings
         test_case = 'call:tag{items:[<|"|>foo<|"|>,<|"|>bar<|"|>]}'
         tool_call = gemma4.parse_tool_call(test_case, None)