fix(proxy): normalize leaked GLM/qwen tool-call arg dialect from responses (#85) (#86)

Some open-weight backends (z-ai/glm-5.2, qwen family) intermittently fail
to decode their own tool-call argument encoding, leaving raw markup inside
the structured `arguments` a well-formed native tool_calls entry returns:
  GLM:  <arg_key>NAME</arg_key> <arg_value>VALUE</arg_value>
  qwen: <parameter=NAME>VALUE</parameter>
The intended payload is intact inside the wrapper, so it's a parse gap not
lost data. Normalize both the value-level leak (dialect inside one value of
an otherwise-valid JSON object, the #85 symptom) and the whole-string leak,
applied to each successful response before it is returned or logged, so no
downstream consumer (client or log) sees the markup. Fully defensive: any
parse failure leaves the value untouched. Repair count recorded as
tool_call_dialect_repaired on the response log so the leak rate stays
measurable. Durable server-side fix for the class geo-agent#276 defends
against client-side.

Author: cboettig

CHANGELOG.md

@@ -8,6 +8,22 @@ See [Releases](README.md#releases) for how a release is cut.
 
 ## [Unreleased]
 
+### Fixed
+- **Strip leaked `<arg_key>`/`<arg_value>` (GLM) and `<parameter=…>` (qwen) tool-call
+  arg dialect from responses (#85).** Some open-weight backends (`z-ai/glm-5.2`, the
+  qwen family) intermittently fail to decode their own tool-call argument encoding,
+  leaving raw markup inside the structured `arguments` a well-formed native
+  `tool_calls` entry returns — e.g. a `value_stats` value arriving as
+  `<arg_key>value_stats</arg_key> <arg_value>{…}</arg_value>` instead of the parsed
+  object. The proxy now normalizes both the value-level leak (dialect inside one value
+  of an otherwise-valid JSON object) and the whole-string leak (the entire `arguments`
+  is raw dialect) in `_normalize_response_tool_calls`, applied to each successful
+  response *before* it is returned or logged, so no downstream consumer (client or log)
+  ever sees the markup. Fully defensive — any parse failure leaves the value untouched.
+  The repair count is recorded as `tool_call_dialect_repaired` on the response log so
+  the leak rate stays measurable. Durable server-side fix for the leak class that
+  geo-agent#276 was defending against client-side.
+
 ### Added
 - **Multi-key client auth — accept more than one `PROXY_KEY` so eval/dev keys are
   independently revocable.** New `PROXY_KEYS_EXTRA` env (comma-separated, wired from

llm_proxy.py

@@ -95,6 +95,98 @@ def _cap(s: Optional[str], limit: int) -> str:
         return s[:limit]
     return s
 
+# --- Tool-call arg-dialect normalization (#85) -------------------------------
+# Some open-weight backends intermittently fail to decode their own tool-call
+# argument encoding, leaving raw markup inside the structured `arguments` string
+# that an otherwise well-formed native `tool_calls` entry hands back:
+#   GLM (z-ai/glm-5.2):  <arg_key>NAME</arg_key> <arg_value>VALUE</arg_value>
+#   qwen / hermes:       <parameter=NAME>VALUE</parameter>
+# The intended payload is intact *inside* the wrapper, so this is a
+# serialization/parse gap, not lost data (#85; qwen precedent geo-agent#276).
+# We repair it here — before the response is returned or logged — so no
+# downstream consumer (client or log) ever sees the dialect. Fully defensive:
+# any parse failure leaves the value untouched.
+_ARG_DIALECT_MARKERS = ("<arg_key>", "<arg_value>", "<parameter=")
+_GLM_ARG_PAIR_RE   = re.compile(r"<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)(?:</arg_value>|$)", re.DOTALL)
+_QWEN_ARG_PAIR_RE  = re.compile(r"<parameter=(.*?)>(.*?)(?:</parameter>|$)", re.DOTALL)
+_GLM_ARG_VALUE_RE  = re.compile(r"<arg_value>(.*?)(?:</arg_value>|$)", re.DOTALL)
+_QWEN_ARG_VALUE_RE = re.compile(r"<parameter=[^>]*>(.*?)(?:</parameter>|$)", re.DOTALL)
+
+
+def _coerce_json(s: str):
+    """Parse `s` as JSON when it is structured data, else return it stripped.
+    Lets an unwrapped `{...}`/`[...]`/number value come back structured while a
+    bare scalar string stays a string."""
+    s = s.strip()
+    try:
+        return json.loads(s)
+    except Exception:
+        return s
+
+
+def _unwrap_dialect_value(val):
+    """If a single tool-call argument *value* carries leaked arg-dialect markup,
+    extract the real payload from inside the <arg_value>/<parameter=…> wrapper.
+    Returns (value, changed)."""
+    if not isinstance(val, str) or not any(m in val for m in _ARG_DIALECT_MARKERS):
+        return val, False
+    m = _GLM_ARG_VALUE_RE.search(val) or _QWEN_ARG_VALUE_RE.search(val)
+    if not m:
+        return val, False
+    return _coerce_json(m.group(1)), True
+
+
+def _normalize_tool_call_arguments(arguments):
+    """Strip leaked arg-dialect markup from a tool call's `arguments` string.
+    Handles both the value-level leak (dialect inside one value of an otherwise
+    valid JSON object, the #85 symptom) and the whole-string leak (the entire
+    `arguments` is raw dialect). Returns (arguments_string, changed)."""
+    if not isinstance(arguments, str) or not any(m in arguments for m in _ARG_DIALECT_MARKERS):
+        return arguments, False
+    # Case 1: arguments is valid JSON; dialect leaked into individual values.
+    try:
+        obj = json.loads(arguments)
+    except Exception:
+        obj = None
+    if isinstance(obj, dict):
+        changed = False
+        for k, v in list(obj.items()):
+            new_v, ch = _unwrap_dialect_value(v)
+            if ch:
+                obj[k] = new_v
+                changed = True
+        return (json.dumps(obj), True) if changed else (arguments, False)
+    # Case 2: the whole arguments string is raw dialect — rebuild the object by
+    # pairing each key tag with the value tag that follows it.
+    pairs = _GLM_ARG_PAIR_RE.findall(arguments) or _QWEN_ARG_PAIR_RE.findall(arguments)
+    if pairs:
+        return json.dumps({k.strip(): _coerce_json(v) for k, v in pairs}), True
+    return arguments, False
+
+
+def _normalize_response_tool_calls(result) -> int:
+    """Repair leaked tool-call arg dialect (#85) in an upstream response, in
+    place. Returns the number of tool-call `arguments` repaired. Fully
+    defensive: any error leaves `result` untouched and returns 0 —
+    normalization must never corrupt a response or break serving."""
+    repaired = 0
+    try:
+        for choice in result.get("choices") or []:
+            message = (choice or {}).get("message") or {}
+            for tc in message.get("tool_calls") or []:
+                fn = (tc or {}).get("function")
+                if not isinstance(fn, dict):
+                    continue
+                new_args, changed = _normalize_tool_call_arguments(fn.get("arguments"))
+                if changed:
+                    fn["arguments"] = new_args
+                    repaired += 1
+    except Exception as e:  # pragma: no cover - defensive
+        print(f"⚠️  tool-call dialect normalization skipped: "
+              f"{type(e).__name__}: {e}", flush=True)
+        return 0
+    return repaired
+
 # --- Credential scrubbing ----------------------------------------------------
 # Credentials reach the logs because the geo-agent `query` MCP tool accepts
 # s3_key/s3_secret in its arguments, which flow through `tool_calls`, tool
@@ -405,7 +497,7 @@ def log_request(provider: str, model: str, messages: List[Dict], tools_count: in
     _emit(log_entry)
 
 @_never_raises
-def log_response(provider: str, model: str, response_data: dict, latency_ms: int, error: str = None, origin: str = None, request_id: str = None, session_id: str = None, client: str = None, upstream_headers: dict = None):
+def log_response(provider: str, model: str, response_data: dict, latency_ms: int, error: str = None, origin: str = None, request_id: str = None, session_id: str = None, client: str = None, upstream_headers: dict = None, dialect_repaired: int = 0):
     """Log response in structured JSON format"""
     log_entry = {
         "timestamp": datetime.utcnow().isoformat() + "Z",
@@ -453,6 +545,12 @@ def log_response(provider: str, model: str, response_data: dict, latency_ms: int
                     for tc in message["tool_calls"]
                 ]
 
+        # How many tool-call arguments were repaired of leaked arg dialect (#85).
+        # Kept queryable so the leak rate stays measurable even though the markup
+        # itself no longer reaches the logs.
+        if dialect_repaired:
+            log_entry["tool_call_dialect_repaired"] = dialect_repaired
+
         # Extract token usage if available
         if "usage" in response_data:
             log_entry["tokens"] = response_data["usage"]
@@ -595,10 +693,18 @@ async def proxy_chat(request: ChatRequest, http_request: Request, authorization:
             response = await http_client.post(endpoint, json=payload, headers=headers)
             response.raise_for_status()
             result = response.json()
-            
+
+            # Repair leaked tool-call arg dialect (#85) before returning OR
+            # logging, so neither the client nor the log ever sees the markup.
+            dialect_repaired = _normalize_response_tool_calls(result)
+            if dialect_repaired:
+                print(f"🧹 Normalized {dialect_repaired} tool-call argument(s) with "
+                      f"leaked arg dialect (model={request.model}, request_id={request_id})",
+                      flush=True)
+
             # Log successful response
             latency_ms = int((time.time() - start_time) * 1000)
-            log_response(provider_name, request.model, result, latency_ms, origin=origin, request_id=request_id, session_id=session_id, client=client)
+            log_response(provider_name, request.model, result, latency_ms, origin=origin, request_id=request_id, session_id=session_id, client=client, dialect_repaired=dialect_repaired)
 
             return result
 

test_logging.py

@@ -498,6 +498,142 @@ def test_openrouter_only_knobs():
     assert "usage" not in nrp_payload, "usage must not leak to non-OpenRouter"
 
 
+# --- Tool-call arg-dialect normalization (#85) -------------------------------
+# glm-5.2 (and the qwen family) intermittently leak their tool-call arg encoding
+# into the structured `arguments`. Verbatim symptom from the issue: a valid outer
+# JSON object whose `value_stats` value is wrapped in the GLM XML arg dialect.
+
+def test_normalize_glm_value_level_leak():
+    """#85: dialect leaked into one value of an otherwise-valid JSON object.
+    The wrapper is stripped and the intended JSON payload comes back structured."""
+    p = importlib.reload(llm_proxy)
+    inner = {"by_res": {"2": {"max": 9.45, "min": 0.1}}}
+    args = json.dumps({
+        "layer_id": "hardwood",
+        "value_stats": f'<arg_key>value_stats</arg_key> <arg_value>{json.dumps(inner)}</arg_value>',
+    })
+    out, changed = p._normalize_tool_call_arguments(args)
+    assert changed
+    parsed = json.loads(out)
+    assert parsed["value_stats"] == inner           # structured, not a string
+    assert parsed["layer_id"] == "hardwood"         # untouched
+    assert "<arg_key>" not in out and "<arg_value>" not in out
+
+
+def test_normalize_glm_value_level_leak_unterminated():
+    """The leaked value may arrive without a closing </arg_value> tag (as the
+    issue's truncated capture showed). We still recover the payload up to end."""
+    p = importlib.reload(llm_proxy)
+    inner = {"by_res": {"2": {"max": 9.45}}}
+    args = json.dumps({
+        "value_stats": f'<arg_key>value_stats</arg_key> <arg_value>{json.dumps(inner)}',
+    })
+    out, changed = p._normalize_tool_call_arguments(args)
+    assert changed
+    assert json.loads(out)["value_stats"] == inner
+
+
+def test_normalize_whole_string_glm_dialect():
+    """The entire `arguments` string is raw GLM dialect (no valid outer JSON)."""
+    p = importlib.reload(llm_proxy)
+    raw = ('<arg_key>layer_id</arg_key> <arg_value>hardwood</arg_value> '
+           '<arg_key>opacity</arg_key> <arg_value>0.5</arg_value>')
+    out, changed = p._normalize_tool_call_arguments(raw)
+    assert changed
+    parsed = json.loads(out)
+    assert parsed == {"layer_id": "hardwood", "opacity": 0.5}
+
+
+def test_normalize_qwen_parameter_dialect():
+    """The qwen/hermes `<parameter=NAME>VALUE</parameter>` form of the same leak."""
+    p = importlib.reload(llm_proxy)
+    args = json.dumps({"sql": "<parameter=sql>SELECT 1</parameter>"})
+    out, changed = p._normalize_tool_call_arguments(args)
+    assert changed
+    assert json.loads(out)["sql"] == "SELECT 1"
+
+
+def test_normalize_leaves_clean_arguments_untouched():
+    """No dialect markers → byte-identical passthrough, no wasted re-serialize."""
+    p = importlib.reload(llm_proxy)
+    args = json.dumps({"sql": "SELECT * FROM t WHERE a < 5", "n": 3})
+    out, changed = p._normalize_tool_call_arguments(args)
+    assert not changed
+    assert out == args
+
+
+def test_normalize_response_tool_calls_in_place_and_counts():
+    """The response-level pass mutates result in place and returns a repair count;
+    a clean sibling tool call in the same response is left alone."""
+    p = importlib.reload(llm_proxy)
+    result = {"choices": [{"message": {"tool_calls": [
+        {"function": {"name": "add_hex_tile_layer", "arguments": json.dumps(
+            {"value_stats": '<arg_value>{"by_res": {"2": {"max": 1}}}</arg_value>'})}},
+        {"function": {"name": "get_schema", "arguments": '{"dataset": "ca"}'}},
+    ]}}]}
+    n = p._normalize_response_tool_calls(result)
+    assert n == 1
+    tcs = result["choices"][0]["message"]["tool_calls"]
+    assert json.loads(tcs[0]["function"]["arguments"])["value_stats"] == {"by_res": {"2": {"max": 1}}}
+    assert tcs[1]["function"]["arguments"] == '{"dataset": "ca"}'
+
+
+def test_normalize_response_is_defensive_on_garbage():
+    """Malformed shapes never raise — normalization must not break serving."""
+    p = importlib.reload(llm_proxy)
+    for junk in ({}, {"choices": None}, {"choices": [None]},
+                 {"choices": [{"message": {"tool_calls": [{"function": None}]}}]},
+                 {"choices": [{"message": {"tool_calls": "nope"}}]}):
+        assert p._normalize_response_tool_calls(junk) == 0
+
+
+def test_handler_repairs_dialect_and_logs_count():
+    """End-to-end: a glm-5.2 response with a leaked value is repaired before it
+    is returned to the client, and the repair count is recorded in the log."""
+    import asyncio
+    from unittest.mock import patch
+
+    p = _reload(PROXY_KEY="testkey")
+    p._log_buffer.clear()
+
+    leaked = json.dumps({"value_stats": '<arg_key>value_stats</arg_key> <arg_value>{"by_res": {"2": {"max": 9.45}}}</arg_value>'})
+
+    class _FakeResp:
+        def raise_for_status(self):
+            pass
+        def json(self):
+            return {"choices": [{"message": {
+                        "content": None,
+                        "tool_calls": [{"function": {"name": "add_hex_tile_layer",
+                                                     "arguments": leaked}}]}}],
+                    "usage": {"total_tokens": 5}}
+
+    class _FakeAsyncClient:
+        def __init__(self, *a, **k):
+            pass
+        async def __aenter__(self):
+            return self
+        async def __aexit__(self, *a):
+            return False
+        async def post(self, *a, **k):
+            return _FakeResp()
+
+    class _FakeRequest:
+        headers = {"origin": "https://ca-30x30.nrp-nautilus.io"}
+
+    req = p.ChatRequest(model="z-ai/glm-5.2", messages=[{"role": "user", "content": "what fraction of ca hardwood is protected?"}])
+    with patch.object(p, "get_provider_for_model",
+                      return_value=("openrouter", {"endpoint": "http://or", "api_key": "k"})), \
+         patch.object(p.httpx, "AsyncClient", _FakeAsyncClient):
+        result = asyncio.run(p.proxy_chat(req, _FakeRequest(), authorization="Bearer testkey"))
+
+    returned = json.loads(result["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"])
+    assert returned["value_stats"] == {"by_res": {"2": {"max": 9.45}}}   # client gets structured data
+    responses = [e for e in p._log_buffer if e.get("type") == "response"]
+    assert responses[0]["tool_call_dialect_repaired"] == 1
+    assert "<arg_key>" not in json.dumps(responses[0])                   # log is clean too
+
+
 if __name__ == "__main__":
     import sys
     fns = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)]