Fix Gemma tool calls losing nested object and array arguments to raw strings

[Adi2K]

Proposed changes

Gemma 4 writes every string in a tool call between escape markers (<|"|>),
including strings inside nested objects and arrays. A call with the
arguments {"location": {"city": "Paris", "country": "FR"}} is written as:

<|tool_call>call:get_weather{location:{city:<|"|>Paris<|"|>,country:<|"|>FR<|"|>}}<tool_call|>

Today GemmaFunctionParser returns the nested value as a string, even when
the tool's schema declares location an object:

location = "{city:<|\"|>Paris<|\"|>,country:<|\"|>FR<|\"|>}"

A tool that validates its input then rejects the call (for example, an MCP
server answers string found, object expected). With this change:

location = {"city": "Paris", "country": "FR"}

More examples of what Gemma writes, and what the parser returns with this
change:

• An array of objects: {stops:[{city:<|"|>Paris<|"|>},{city:<|"|>Lyon<|"|>}]}
  becomes [{"city": "Paris"}, {"city": "Lyon"}] (today: a string).
• A string with a line break: {note:{text:<|"|>line one + newline +
  line two<|"|>}} becomes {"text": "line one\nline two"} (today: a string).
• A flat value such as {city:<|"|>Paris<|"|>} already worked and is
  unchanged.

Why. #557 added BareKeyJSONParser, which quotes bare object keys and
retries JSON parsing. It only knows JSON's own double quotes, so a nested
value that holds marker-quoted strings still falls back to the raw text. The
#557 review noted this: nested values written as <|"|>…<|"|> are something
"this parser still does not rewrite, but that is probably a separate
follow-up." This PR is that follow-up.

This is the shape the model is trained to write. The chat template's
format_argument macro wraps every string value in <|"|> at any nesting
depth and inserts the text as is, so line breaks can appear inside
(google/gemma-4-e2b-it chat_template.jinja, lines 124–155; tool calls use
it at line 253). The examples above are rendered by that template. Google's
tokenizer_config.json for the same model tells transformers to read these
strings (string_delims: <|"|>, unquoted_keys: true), and mlx-lm's
gemma4 tool parser reads them too. The regex in Google's function-calling
docs is a flat demo that does not read nested values, so it is not a useful
reference here.

What changed. GemmaFunctionParser still tries the original text first,
exactly as before. Only if that fails, and the text contains the marker, it
rewrites each <|"|>…<|"|> span as a JSON string and tries again. The
rewrite escapes backslash, double quote, and every control character below
U+0020. Because the original text is always tried first, anything that
parsed before gives the same result now. This matters for FunctionGemma,
which shares this parser with the <escape> marker:
{a:"<escape>hi<escape>"} still parses as {"a": "<escape>hi<escape>"}.

Two key shapes that the template never writes now parse as well, nested and
at the top level: a marker-quoted key, as in
{<|"|>city<|"|>:<|"|>Paris<|"|>}, and a space before a key, as in
{city:<|"|>Paris<|"|>, country:<|"|>FR<|"|>}. gemma-4-e2b-it-4bit wrote
both shapes inside a nested value with thinking disabled: the quoted key in
1 of 8 sampled runs, the space in 2 of 8. Nested, the rewrite above turns a
marker-quoted key into a JSON string, which BareKeyJSONParser accepts. A
top-level key is read directly, so the parser now trims whitespace around
it and strips a marker pair around it, which transformers and mlx-lm already
do at every depth because they rewrite the whole argument text. Before,
{city:<|"|>Paris<|"|>, days:3}
stored days under the key " days" and {<|"|>city<|"|>:<|"|>Paris<|"|>}
stored city under <|"|>city<|"|>, neither of which the tool's schema
knows.

How it's tested. Tests/MLXLMTests/ToolTests.swift adds 22 tests next to
the existing Gemma tests, five of them parameterized.

19 fail without these changes (31 issues): nested objects, arrays of objects, bare
scalars mixed with marker strings, line breaks, tabs and other control
characters, quotes and backslashes, empty and backslash-ended strings,
non-ASCII text, marker-quoted keys, the two key shapes above, a space before
a top-level key, a marker-quoted top-level key, a marker
string in an array nested several objects deep, a call with no tool schema,
a FunctionGemma nested <escape> string, a FunctionGemma value where an
unpaired marker follows a paired one, and a streamed call through
ToolCallProcessor(format: .gemma4).

3 pass both before and after. They use input outside the dialect and pin
existing behavior: JSON-quoted FunctionGemma strings that contain markers
(the original text is read first), a FunctionGemma value with an unpaired
marker, and a string-typed parameter that keeps its raw brace text. The test
with an unpaired marker after a paired one also uses malformed input; it
pins the result this change chooses.

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:

  I used Claude (via Claude Code) for this change. It wrote the parser change and the tests,
  ran the unit tests with and without the fix, ran pre-commit run --all-files and scripts/verify-docs.sh, and drafted this description
  and the commit message. Separate AI review passes checked the change and
  found a FunctionGemma regression in an earlier version, which is fixed
  here. A later AI cross-check rendered the Gemma 4 chat template, compared
  other Gemma 4 parsers, and captured gemma-4-e2b-it output to confirm which
  shapes the model writes. I have read the diff, understand it, and take
  responsibility for it.

[aleroot]

This trims the key but doesn't unmark it, so a marker-quoted top-level key still lands in arguments verbatim.

testGemmaMarkerQuotedNestedKeys establishes that the model really does emit marker-quoted keys this PR fixes that shape one level down via quotingMarkedStrings , but the same shape at the top level is still broken...

[Adi2K]

Right — the trim covered the space but not the markers, so {<|"|>city<|"|>:<|"|>Paris<|"|>} stored the key as <|"|>city<|"|>. Nested keys were already fine because the rewrite turns them into JSON strings and BareKeyJSONParser takes a quoted key; the top-level key is read directly. Pushed a commit that strips a marker pair around the trimmed key and adds a test for it next to the space case. I have only seen the model do this nested, not at the top level, but it is the same shape, and transformers and mlx-lm read it at every depth.

[aleroot]

The first structuredValues.parse(literal) is dead work whenever the literal contains a marker.
BareKeyJSONParser can never succeed on such input: <|"|> carries an unescaped " , so tryParseJSON rejects it, and quotingBareKeys passes values through verbatim
and rejects marker-wrapped keys ( quotedKey fails on < ), so the retry string still contains the marker. Checking contains(marker) first is therefore behaviour-preserving, not just an optimisation:

 private func parseStructured(_ literal: String, marker: String) -> (any Sendable)? { 
      guard literal.contains(marker) else { return structuredValues.parse(literal) } 
      return structuredValues.parse(Self.quotingMarkedStrings(literal, marker: marker)) 
  } 

[Adi2K]

For <|"|> that's true: its " is never valid JSON in place, so the first parse can't succeed and the guard changes nothing.

It isn't true for FunctionGemma's <escape>, which has no quote in it and can sit inside a JSON-quoted string. {a:"<escape>hi<escape>"} parses on main today through the bare-key retry as {"a": "<escape>hi<escape>"}. With the guard, the rewrite runs first and produces {a:""hi""}, which nothing parses, so the value falls back to the raw string. The two cases in testFunctionGemmaMarkersInsideQuotedString fail with the guard applied.

So I kept the order and added a doc comment on parseStructured saying why the original text goes first. The cost is one failed JSONSerialization on a nested literal that contains <|"|>.

[aleroot]

Good test for the trim. When my other comment gets addressed, this is the natural home for a marker-quoted top-level key case.

[Adi2K]

Added testGemmaMarkerQuotedTopLevelKey right after it: a marker-quoted key, and a marker-quoted key after a space, so the trim-then-unmark order is pinned too.

[aleroot]

Thank you for raising this pull request, this was something I had in my TODO list.

[Adi2K]

@aleroot one idea for after this lands, not for this PR. The "which parse goes first" question in parseStructured only exists because the dialect gets rewritten into JSON text for JSONSerialization. Your PythonLiteralParser from #531 reads its dialect directly, and the same approach fits here: a small recursive-descent reader for marker strings, JSON strings, bare keys and nested values. Top-level and nested keys would share one path, and " versus the marker would be settled by position. It would replace BareKeyJSONParser and the marker rewrite, with the tests from this PR as the parity gate. Would you take that as a follow-up? I can write up the grammar first if you'd rather see it before code.

[aleroot]

@aleroot one idea for after this lands, not for this PR. The "which parse goes first" question in parseStructured only exists because the dialect gets rewritten into JSON text for JSONSerialization. Your PythonLiteralParser from #531 reads its dialect directly, and the same approach fits here: a small recursive-descent reader for marker strings, JSON strings, bare keys and nested values. Top-level and nested keys would share one path, and " versus the marker would be settled by position. It would replace BareKeyJSONParser and the marker rewrite, with the tests from this PR as the parity gate. Would you take that as a follow-up? I can write up the grammar first if you'd rather see it before code.

@Adi2K yes, it sounds good. If you can open the issue so that I will definitely remember next week, as this week I have no time for this project. Thanks.

[Adi2K]

Opened #628 with the grammar. To be clear, I'll write it myself once this lands. It only needs a review from you, no work on your side.

[aleroot]

Opened #628 with the grammar. To be clear, I'll write it myself once this lands. It only needs a review from you, no work on your side.

Perfetto.