[Frontend] Prevent Kimi K3 reserved markers in response text (#52889)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: OpenAI Codex <codex@openai.com>

rust/src/parser/src/unified/kimi_k3.rs

@@ -56,6 +56,7 @@ use crate::unified::parsing_failed;
 use crate::utils::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker};
 
 const OPEN: &str = "<|open|>";
+const CLOSE: &str = "<|close|>";
 const SEP: &str = "<|sep|>";
 const END_OF_MSG: &str = "<|end_of_msg|>";
 

rust/src/parser/src/unified/kimi_k3/structural_tag.rs

@@ -12,8 +12,8 @@ use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, StructuralTag, T
 use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam, function_parameters};
 
 use super::{
-    ARG_CLOSE, CALL_CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN, RESPONSE_CLOSE,
-    RESPONSE_OPEN, SEP, THINK_CLOSE, TOOLS_CLOSE, TOOLS_OPEN,
+    ARG_CLOSE, CALL_CLOSE, CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN,
+    RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, TOOLS_CLOSE, TOOLS_OPEN,
 };
 
 pub(super) static KIMI_K3_STRUCTURAL_TAG_BUILDER: KimiK3StructuralTagBuilder =
@@ -64,7 +64,9 @@ fn response_prefix(reasoning: bool) -> Vec<Format> {
     }
     elements.push(Format::tag(
         "",
-        Format::any_text_excluding(&[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]),
+        // Keep marker-looking prefixes out of response text so the grammar
+        // must resolve them through a valid channel boundary.
+        Format::any_text_excluding(&[OPEN, CLOSE, END_OF_MSG]),
         RESPONSE_CLOSE,
     ));
     elements
@@ -384,7 +386,7 @@ mod tests {
         )
         .unwrap();
 
-        expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|close|>response<|sep|>","<|open|>tools<|sep|>","<|close|>message<|sep|>","<|end_of_msg|>"]},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap());
+        expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|open|>","<|close|>","<|end_of_msg|>"]},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap());
     }
 
     #[test]

tests/tool_parsers/test_structural_tag_registry.py

@@ -177,6 +177,7 @@ _K3_TOOLS_CLOSE = "<|close|>tools<|sep|>"
 _K3_CALL_CLOSE = "<|close|>call<|sep|>"
 _K3_ARG_CLOSE = "<|close|>argument<|sep|>"
 _K3_MESSAGE_CLOSE = "<|close|>message<|sep|>"
+_K3_END_OF_MSG = "<|end_of_msg|>"
 
 
 def _k3_tools_by_name() -> list[ChatCompletionToolsParam]:
@@ -357,6 +358,20 @@ def test_kimi_k3_auto_strict_allows_response_only(sample_tools_strict):
     assert _is_grammar_accept_string(grammar, _k3_response("Just answering."))
 
 
+@pytest.mark.parametrize(
+    "body",
+    [
+        _K3_RESPONSE_OPEN + "answer<|close|>message",
+        _K3_RESPONSE_OPEN + "answer<|open|>tools",
+        _K3_RESPONSE_OPEN + "answer" + _K3_END_OF_MSG,
+    ],
+)
+def test_kimi_k3_response_body_rejects_reserved_marker_prefixes(body: str):
+    assert not _is_grammar_accept_string(
+        _k3_grammar("required"), body, require_termination=False
+    )
+
+
 @pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS))
 def test_get_model_structural_tag_supports_named_tool_choice(
     model: str,

vllm/tool_parsers/structural_tag_registry.py

@@ -376,6 +376,7 @@ _K3_ARG_CLOSE = f"{_K3_CLOSE}argument{_K3_SEP}"
 # the end-of-message token. It is generated (not part of the prompt prefix), so
 # the tag must permit it or the FSM would mask the model's natural terminator.
 _K3_MESSAGE_CLOSE = f"{_K3_CLOSE}message{_K3_SEP}"
+_K3_END_OF_MSG = "<|end_of_msg|>"
 
 # JSON-schema type -> K3 XTML ``type=`` attribute value. Mirrors
 # ``encoding_k3._xtml_type`` (integer collapses onto number).
@@ -556,12 +557,17 @@ def _k3_response_prefix() -> list[Any]:
 
     ``response`` is generated in thinking mode (prefix ends at
     ``<|open|>think<|sep|>``) but is part of the generation prefix in
-    non-thinking mode, so its open marker is optional. The body is bounded by
-    the response close marker.
+    non-thinking mode, so its open marker is optional. Excluding the reserved
+    marker heads commits ``<|close|>`` to the response-close branch immediately
+    and prevents a nested channel from starting inside response text.
     """
     return [
         OptionalFormat(content=ConstStringFormat(value=_K3_RESPONSE_OPEN)),
-        TagFormat(begin="", content=AnyTextFormat(), end=_K3_RESPONSE_CLOSE),
+        TagFormat(
+            begin="",
+            content=AnyTextFormat(excludes=[_K3_OPEN, _K3_CLOSE, _K3_END_OF_MSG]),
+            end=_K3_RESPONSE_CLOSE,
+        ),
     ]