text: fix streaming tool-call spill; gate native markup unconditionally; v0.1.2

The streaming content gate only ran when the request carried a tools array. But
the AISBF relay can drop tools (as it does max_tokens), so a gemma model still
emitting <|tool_call>call:NAME{…} from its system prompt had that markup streamed
straight to the client as content — the spill that poisoned history and fed the
tool-call loop.

_gate_tool_content now:
  * ALWAYS withholds the unambiguous native special-token markers (<|tool_call>,
    <|tool_response>, DeepSeek DSML) and their cross-chunk partials, even with no
    tools declared;
  * gates the ambiguous <tool>…</tool> XML form only when tools are present;
  * gates the gemma-4 call:NAME{…} form by the per-model gemma_tool_parser mode
    (off: never; full: always; restricted: only declared tool names), so legit
    call:foo{…} prose/code is streamed instead of withheld+dropped.
The gate is now invoked for every stream (not just tools-present), and the final
flush likewise, with the mode/tool context resolved once per request.

Bump version to 0.1.2.
Co-Authored-By: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RdMufYvtTbtGDWsiZVoXce
parent c7105023
...@@ -16,7 +16,7 @@ ...@@ -16,7 +16,7 @@
# Canonical product version for CoderAI — single source of truth. Both the API # Canonical product version for CoderAI — single source of truth. Both the API
# metadata and the admin web UI read from here. # metadata and the admin web UI read from here.
__version__ = "0.1.1" __version__ = "0.1.2"
# Configure the CUDA caching allocator BEFORE torch is imported anywhere. # Configure the CUDA caching allocator BEFORE torch is imported anywhere.
# expandable_segments lets the allocator return freed pages to the driver even # expandable_segments lets the allocator return freed pages to the driver even
......
...@@ -2269,6 +2269,10 @@ _TOOL_SPAN_RE = _re.compile(r'<(tool|tool_call)\b[\s\S]*?</\1\s*>', _re.IGNORECA ...@@ -2269,6 +2269,10 @@ _TOOL_SPAN_RE = _re.compile(r'<(tool|tool_call)\b[\s\S]*?</\1\s*>', _re.IGNORECA
_TOOL_OPEN_RE = _re.compile(r'<(?:tool|tool_call)\b', _re.IGNORECASE) _TOOL_OPEN_RE = _re.compile(r'<(?:tool|tool_call)\b', _re.IGNORECASE)
_TOOL_OPEN_TAGS = ('<tool>', '<tool_call>', '<|tool_call>', '<|tool_call|>', _TOOL_OPEN_TAGS = ('<tool>', '<tool_call>', '<|tool_call>', '<|tool_call|>',
'<|dsml|tool_calls>') '<|dsml|tool_calls>')
# Native/special-token tool markers used for trailing-partial holding even when no
# tools were declared (so a relay-stripped tools array can't let <|tool_call> leak).
_NATIVE_PARTIAL_TAGS = ('<|tool_call>', '<|tool_call|>', '<|tool_response>',
'<|dsml|', '<|dsml|')
# gemma/qwen native special-token tool marker `<|tool_call>` — usually a special # gemma/qwen native special-token tool marker `<|tool_call>` — usually a special
# token stripped on decode, but some GGUFs emit it as plain text. Treat it like a # token stripped on decode, but some GGUFs emit it as plain text. Treat it like a
# tool-open: withhold everything from it to the end so the raw marker (and the # tool-open: withhold everything from it to the end so the raw marker (and the
...@@ -2319,62 +2323,81 @@ def _gate_reasoning(buffer: str, final: bool = False): ...@@ -2319,62 +2323,81 @@ def _gate_reasoning(buffer: str, final: bool = False):
return buffer, "", "", False return buffer, "", "", False
def _gate_tool_content(buffer: str, final: bool = False): def _gate_tool_content(buffer: str, final: bool = False,
gemma_mode: str = "full", tool_names=None,
gate_xml: bool = True):
"""Split accumulated stream text into (content_to_emit, held_buffer). """Split accumulated stream text into (content_to_emit, held_buffer).
During tool-enabled streaming the model emits ``<tool>{json}</tool>`` spans Tool-call markup must NOT reach the client as visible ``content`` (it's
inline. Those must NOT reach the client as visible ``content`` (they're
surfaced separately as structured ``tool_calls``); otherwise the raw tags leak surfaced separately as structured ``tool_calls``); otherwise the raw tags leak
into the chat. This withholds any complete or in-progress tool span, plus a into the chat — the "spill". The UNAMBIGUOUS native special-token markers
trailing partial ``<`` that could still grow into a tool tag, and streams only (``<|tool_call>``, ``<|tool_response>``, DeepSeek DSML) are ALWAYS withheld,
the safe text around them. With ``final=True`` any leftover (possibly unclosed) even when no tools were declared on the request — the AISBF relay can drop the
tool span is dropped and the rest emitted. ``tools`` array, but a gemma/qwen model still emits those markers from its
system prompt, and they are never legitimate content.
``gate_xml`` (set from "tools present") controls the ambiguous ``<tool>…</tool>``
XML form. The gemma-4 ``call:NAME{…}`` form is gated by ``gemma_mode``: 'off'
never withholds it; 'full' always; 'restricted' only when NAME is a declared
tool (so legit ``call:foo{…}`` prose/code is streamed, not eaten). With
``final=True`` any leftover (possibly unclosed) withheld span is dropped.
""" """
emit = [] emit = []
# Strip complete tool spans, emitting the text around each. names = set(tool_names) if tool_names else set()
while True: # Complete <tool>…</tool> spans (ambiguous XML form — only when tools present).
m = _TOOL_SPAN_RE.search(buffer) if gate_xml:
if not m: while True:
break m = _TOOL_SPAN_RE.search(buffer)
emit.append(buffer[:m.start()]) if not m:
buffer = buffer[m.end():] break
# An open tag with no close yet → hold from there (a call is in progress). emit.append(buffer[:m.start()])
m = _TOOL_OPEN_RE.search(buffer) buffer = buffer[m.end():]
if m: # Native <|tool_call …> / <|tool_response …> — ALWAYS withhold to the end.
emit.append(buffer[:m.start()])
held = '' if final else buffer[m.start():]
return ''.join(emit), held
# gemma/qwen native `<|tool_call>` marker — withhold from it to the end.
nm = _NATIVE_TOOL_OPEN_RE.search(buffer) nm = _NATIVE_TOOL_OPEN_RE.search(buffer)
if nm: if nm:
emit.append(buffer[:nm.start()]) emit.append(buffer[:nm.start()])
held = '' if final else buffer[nm.start():] return ''.join(emit), ('' if final else buffer[nm.start():])
return ''.join(emit), held # DeepSeek V4 DSML tool call (<|DSML|…>) — ALWAYS withhold to the end.
# gemma-4 `call:NAME{…}` — withhold from the call onward (extracted at the end).
gm = _GEMMA_CALL_OPEN_RE.search(buffer)
if gm:
emit.append(buffer[:gm.start()])
held = '' if final else buffer[gm.start():]
return ''.join(emit), held
# DeepSeek V4 DSML tool call (<|DSML|…>) — withhold from the marker onward.
dm = _DSML_OPEN_RE.search(buffer) dm = _DSML_OPEN_RE.search(buffer)
if dm: if dm:
emit.append(buffer[:dm.start()]) emit.append(buffer[:dm.start()])
held = '' if final else buffer[dm.start():] return ''.join(emit), ('' if final else buffer[dm.start():])
return ''.join(emit), held # Ambiguous XML open tag with no close yet (only when tools present).
# Hold back a trailing '<…' that could still become a tool open tag. if gate_xml:
m = _TOOL_OPEN_RE.search(buffer)
if m:
emit.append(buffer[:m.start()])
return ''.join(emit), ('' if final else buffer[m.start():])
# gemma-4 `call:NAME{…}` — withhold per mode. Under 'restricted', a name that
# isn't a declared tool is left in place (legit content) and scanning continues.
if gemma_mode != "off":
pos = 0
while True:
gm = _GEMMA_CALL_OPEN_RE.search(buffer, pos)
if not gm:
break
_nm = _re.match(r'call:\s*([A-Za-z_]\w*)', buffer[gm.start():])
cname = _nm.group(1) if _nm else ''
if gemma_mode == "restricted" and cname not in names:
pos = gm.end()
continue
emit.append(buffer[:gm.start()])
return ''.join(emit), ('' if final else buffer[gm.start():])
# Hold back trailing partials that could still grow into a marker next chunk.
if not final: if not final:
lt = buffer.rfind('<') lt = buffer.rfind('<')
if lt != -1: if lt != -1:
tail = buffer[lt:].lower() tail = buffer[lt:].lower()
if any(t.startswith(tail) for t in _TOOL_OPEN_TAGS): # Native/DSML partials always; XML partials only when gating XML.
cand = _TOOL_OPEN_TAGS if gate_xml else _NATIVE_PARTIAL_TAGS
if any(t.startswith(tail) for t in cand):
emit.append(buffer[:lt]) emit.append(buffer[:lt])
return ''.join(emit), buffer[lt:] return ''.join(emit), buffer[lt:]
# Hold a trailing 'call:NAME' (no '{' yet) that may grow into a gemma call. if gemma_mode != "off":
cm = _re.search(r'call:\s*[A-Za-z_]?\w*$', buffer) cm = _re.search(r'call:\s*[A-Za-z_]?\w*$', buffer)
if cm: if cm:
emit.append(buffer[:cm.start()]) emit.append(buffer[:cm.start()])
return ''.join(emit), buffer[cm.start():] return ''.join(emit), buffer[cm.start():]
emit.append(buffer) emit.append(buffer)
return ''.join(emit), '' return ''.join(emit), ''
...@@ -2522,6 +2545,23 @@ async def stream_chat_response( ...@@ -2522,6 +2545,23 @@ async def stream_chat_response(
model=model_name or "", task_id=request_id) model=model_name or "", task_id=request_id)
task_registry.start(_tid) task_registry.start(_tid)
# Tool-content gating context (resolved once). Native markers are always gated;
# XML <tool> gating follows "tools present"; the gemma call:NAME{ form follows
# the per-model gemma_tool_parser mode. Resolving here lets the gate withhold
# native markup even when the relay dropped the tools array (the spill).
from codai.models.parser import resolve_gemma_tool_mode as _resolve_gemma_mode
_gemma_mode = _resolve_gemma_mode(model_name)
_gate_xml = bool(tools)
_tool_names = set()
try:
for _t in (tools or []):
_fn = _t.get("function") if isinstance(_t, dict) else getattr(_t, "function", None)
_n = (_fn.get("name") if isinstance(_fn, dict) else getattr(_fn, "name", None)) if _fn else None
if _n:
_tool_names.add(_n)
except Exception:
_tool_names = set()
try: try:
chunk_count = 0 chunk_count = 0
_gen_t0 = None # wall-clock of the first generated token (for it/s) _gen_t0 = None # wall-clock of the first generated token (for it/s)
...@@ -2665,15 +2705,19 @@ async def stream_chat_response( ...@@ -2665,15 +2705,19 @@ async def stream_chat_response(
await asyncio.sleep(0) await asyncio.sleep(0)
continue continue
# When tools are enabled, gate the content so in-progress <tool>…</tool> # Gate the content so tool-call markup is never streamed as visible text
# spans are never streamed as visible text (they're surfaced as # (it's surfaced as structured tool_calls after the stream). Native
# structured tool_calls after the stream). Without tools, stream as-is. # markers (<|tool_call>, DSML) are withheld ALWAYS — even when no tools
if tools: # were declared, since the relay can drop the tools array yet the model
content_buffer += filtered_chunk # still emits them (the spill). XML <tool> and gemma call:NAME{ gating
filtered_chunk, content_buffer = _gate_tool_content(content_buffer) # follow tools-present / the gemma mode.
if not filtered_chunk: content_buffer += filtered_chunk
await asyncio.sleep(0) filtered_chunk, content_buffer = _gate_tool_content(
continue content_buffer, gemma_mode=_gemma_mode,
tool_names=_tool_names, gate_xml=_gate_xml)
if not filtered_chunk:
await asyncio.sleep(0)
continue
data = { data = {
"id": completion_id, "id": completion_id,
...@@ -2721,8 +2765,11 @@ async def stream_chat_response( ...@@ -2721,8 +2765,11 @@ async def stream_chat_response(
# Flush any safe trailing text held back by the tool-content gate # Flush any safe trailing text held back by the tool-content gate
# (dropping leftover/unclosed tool tags — they become tool_calls below). # (dropping leftover/unclosed tool tags — they become tool_calls below).
if tools and content_buffer: # Runs even without a tools array, since native markup is gated regardless.
tail_content, content_buffer = _gate_tool_content(content_buffer, final=True) if content_buffer:
tail_content, content_buffer = _gate_tool_content(
content_buffer, final=True, gemma_mode=_gemma_mode,
tool_names=_tool_names, gate_xml=_gate_xml)
if tail_content: if tail_content:
data = { data = {
"id": completion_id, "id": completion_id,
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment