# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for the new declarative response_template parser.

All six real-model template fixtures from the legacy test suite are re-expressed
here in the new region-spec shape and asserted against the same expected
output dicts. Any divergence indicates a regression in the new executor."""

import random
import tempfile
import unittest

from transformers import AutoTokenizer
from transformers.utils.chat_parsing import ResponseParser, parse_response


cohere_template = {
    "defaults": {"role": "assistant"},
    "start_anchor": "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>",
    "fields": {
        "content": {
            "open": "<|START_RESPONSE|>",
            "close": "<|END_RESPONSE|>",
            "content": "text",
        },
        "thinking": {
            "open": "<|START_THINKING|>",
            "close": "<|END_THINKING|>",
            "content": "text",
        },
        "tool_calls": {
            "open": "<|START_ACTION|>",
            "close": "<|END_ACTION|>",
            "content": "json",
            "transform_each": True,
            "transform": {"type": "function", "function": {"name": "{tool_name}", "arguments": "{parameters}"}},
        },
    },
}


ernie_template = {
    "defaults": {"role": "assistant"},
    "start_anchor": "Assistant:",
    "fields": {
        "thinking": {
            "open_pattern": r"(?:^|<think>\s*)",
            "close": "</think>",
            "content": "text",
        },
        "content": {
            "open": "<response>\n",
            "close_pattern": r"\n?</response>",
            "content": "text",
        },
        "tool_calls": {
            "open": "<tool_call>",
            "close": "</tool_call>",
            "repeats": True,
            "content": "json",
            "transform": {"type": "function", "function": "{content}"},
        },
    },
}


gpt_oss_template = {
    "defaults": {"role": "assistant"},
    "start_anchor": "<|start|>assistant",
    "fields": {
        "thinking": {
            "open": "<|channel|>analysis<|message|>",
            "close": "<|end|>",
            "content": "text",
        },
        "content": {
            "open": "<|channel|>final<|message|>",
            "close": "<|end|>",
            "content": "text",
        },
        "tool_calls": {
            "open_pattern": r"<\|channel\|>commentary to=functions\.(?P<name>\w+).*?<\|message\|>",
            "close": "<|call|>",
            "repeats": True,
            "content": "json",
            "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}},
        },
    },
}


smollm_template = {
    "defaults": {"role": "assistant"},
    "start_anchor": "<|im_start|>assistant\n",
    "fields": {
        "thinking": {"open": "<think>", "close": "</think>", "content": "text"},
        "tool_calls": {
            "open": "<tool_call>",
            "close": "</tool_call>",
            "repeats": True,
            "content": "json",
            "transform": {"type": "function", "function": "{content}"},
        },
        "content": {
            "close": "<|im_end|>",
            "content": "text",
        },
    },
}


qwen3_template = {
    "defaults": {"role": "assistant"},
    "start_anchor": "<|im_start|>assistant\n",
    "fields": {
        "thinking": {"open": "<think>", "close": "</think>", "content": "text"},
        "tool_calls": {
            "open_pattern": r"<tool_call>\s*<function=(?P<name>\w+)>",
            "close": "</tool_call>",
            "repeats": True,
            "content": "xml-inline",
            "content_args": {
                "tag_pattern": r"<parameter=(?P<key>\w+)>\s*(?P<value>.*?)\s*</parameter>",
                "value_parser": {"name": "json", "args": {"allow_non_json": True}},
            },
            "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}},
        },
    },
}


gemma4_template = {
    "defaults": {"role": "assistant"},
    # The chat template only emits `<|turn>model\n` when the previous message wasn't a tool_call/
    # tool_response. After a tool_response the prefix just ends with `<tool_response|>` and the
    # model continues from there, so we accept either anchor and truncate past the latest one.
    "start_anchor": ["<|turn>model\n", "<tool_response|>"],
    "fields": {
        "thinking": {
            "open": "<|channel>thought\n",
            "close": "<channel|>",
            "content": "text",
        },
        "tool_calls": {
            "open_pattern": r"<\|tool_call>call:(?P<name>\w+)",
            "close": "<tool_call|>",
            "repeats": True,
            "content": "json",
            "content_args": {
                "unquoted_keys": True,
                "string_delims": [['<|"|>', '<|"|>']],
            },
            "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}},
        },
        "content": {
            "close": ["<turn|>", "<|tool_response>", "<eos>"],
            "content": "text",
        },
    },
}


class ChatResponseTemplateParserTest(unittest.TestCase):
    def test_response_template_save_load(self):
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        tokenizer.response_template = ernie_template
        with tempfile.TemporaryDirectory() as tmpdir:
            tokenizer.save_pretrained(tmpdir)
            reloaded = AutoTokenizer.from_pretrained(tmpdir)
        self.assertEqual(reloaded.response_template, ernie_template)

    def test_tokenizer_parse_response(self):
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        tokenizer.response_template = cohere_template
        model_out = (
            "<|START_THINKING|>I should call a tool.<|END_THINKING|>"
            '<|START_ACTION|>[\n    {"tool_call_id": "0", "tool_name": "simple_tool", '
            '"parameters": {"temperature_format": "Celsius"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>'
        )
        expected = {
            "role": "assistant",
            "thinking": "I should call a tool.",
            "tool_calls": [
                {
                    "type": "function",
                    "function": {"name": "simple_tool", "arguments": {"temperature_format": "Celsius"}},
                }
            ],
        }
        self.assertEqual(tokenizer.parse_response(model_out, prefix=""), expected)

    def test_token_id_inputs(self):
        tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
        tokenizer.response_template = cohere_template
        model_out = (
            "<|START_THINKING|>I should call a tool.<|END_THINKING|>"
            '<|START_ACTION|>[\n    {"tool_call_id": "0", "tool_name": "simple_tool", '
            '"parameters": {"temperature_format": "Celsius"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>'
        )
        parsed = tokenizer.parse_response(model_out, prefix="")
        tokenized = tokenizer(model_out).input_ids
        self.assertEqual(tokenizer.parse_response(tokenized, prefix=""), parsed)

    def test_batched_response(self):
        """A batch of responses (list of strings or list of token-id sequences) returns one parsed
        dict per item; a single-item batch still returns a one-element list, not a bare dict."""
        tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
        tokenizer.response_template = cohere_template
        out_a = (
            "<|START_THINKING|>Think A.<|END_THINKING|>"
            '<|START_ACTION|>[\n    {"tool_call_id": "0", "tool_name": "tool_a", '
            '"parameters": {"x": "1"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>'
        )
        out_b = (
            "<|START_THINKING|>Think B.<|END_THINKING|>"
            '<|START_ACTION|>[\n    {"tool_call_id": "0", "tool_name": "tool_b", '
            '"parameters": {"y": "2"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>'
        )
        single_a = tokenizer.parse_response(out_a, prefix="")
        single_b = tokenizer.parse_response(out_b, prefix="")
        self.assertNotEqual(single_a, single_b)
        # A list of strings is parsed as a batch, one dict per item.
        self.assertEqual(tokenizer.parse_response([out_a, out_b], prefix=""), [single_a, single_b])
        # Batched token-id input (list of token-id sequences) parses the same way.
        ids = [tokenizer(out_a).input_ids, tokenizer(out_b).input_ids]
        self.assertEqual(tokenizer.parse_response(ids, prefix=""), [single_a, single_b])
        # A single-item batch returns a one-element list, not a bare dict.
        self.assertEqual(tokenizer.parse_response([out_a], prefix=""), [single_a])

    def test_explicit_template_schema_detection(self):
        """An explicit new-style template passed as `schema=` is routed to the response-template
        parser, not the legacy `response_schema` parser. New-style is identified by a top-level
        `version` key (the canonical marker) or a `fields` key for templates that omit it."""
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        model_out = (
            "<|START_THINKING|>I should call a tool.<|END_THINKING|>"
            '<|START_ACTION|>[\n    {"tool_call_id": "0", "tool_name": "simple_tool", '
            '"parameters": {"temperature_format": "Celsius"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>'
        )
        expected = parse_response(model_out, cohere_template, prefix="")
        # Detected via the canonical `version` marker...
        self.assertEqual(
            tokenizer.parse_response(model_out, schema={"version": 1, **cohere_template}, prefix=""), expected
        )
        # ...and via `fields` when the template omits `version`.
        self.assertEqual(tokenizer.parse_response(model_out, schema=cohere_template, prefix=""), expected)

    def test_cohere(self):
        model_out = (
            "<|START_THINKING|>I should call a tool.<|END_THINKING|>"
            '<|START_ACTION|>[\n    {"tool_call_id": "0", "tool_name": "simple_tool", '
            '"parameters": {"temperature_format": "Celsius"}}\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>'
        )
        self.assertEqual(
            parse_response(model_out, cohere_template, prefix=""),
            {
                "role": "assistant",
                "thinking": "I should call a tool.",
                "tool_calls": [
                    {
                        "type": "function",
                        "function": {"name": "simple_tool", "arguments": {"temperature_format": "Celsius"}},
                    }
                ],
            },
        )

    def test_ernie_with_tools(self):
        model_out = (
            "The user is asking about the weather in Paris today. Let me check the available tools. "
            "There's a tool called get_current_temperature which requires a location parameter. Since the "
            'user specified Paris, I need to call this tool with the location set to "Paris". I should '
            "make sure the argument is correctly formatted as a string. No other tools are available, so "
            "this is the right one to use. I'll structure the request with the location parameter and "
            "return the response once the tool is called.\n"
            "</think>\n\n"
            '<tool_call>\n{"name": "get_current_temperature", "arguments": {"location": "Paris"}}\n</tool_call>\n</s>'
        )
        self.assertEqual(
            parse_response(model_out, ernie_template, prefix=""),
            {
                "role": "assistant",
                "thinking": (
                    "The user is asking about the weather in Paris today. Let me check the available tools. "
                    "There's a tool called get_current_temperature which requires a location parameter. Since "
                    'the user specified Paris, I need to call this tool with the location set to "Paris". I '
                    "should make sure the argument is correctly formatted as a string. No other tools are "
                    "available, so this is the right one to use. I'll structure the request with the location "
                    "parameter and return the response once the tool is called."
                ),
                "tool_calls": [
                    {
                        "type": "function",
                        "function": {"name": "get_current_temperature", "arguments": {"location": "Paris"}},
                    }
                ],
            },
        )

    def test_ernie_no_tools(self):
        model_out = (
            'The user just greeted me with "Hi! How are you?" I need to respond in a friendly and helpful '
            "manner. Let me start by acknowledging their greeting. I should ask them how they're doing to "
            "engage in conversation.\n\n"
            "First, I'll say hello back and then ask how they're feeling. It's important to show genuine "
            "interest. Maybe mention that I'm here to help with anything they need. Keep the tone warm and "
            "positive. Let me make sure the response is concise but friendly. Alright, that should work.\n"
            "</think>\n\n"
            "<response>\nHello! I'm doing well, thank you for asking. How about you? Is there something "
            "specific you'd like help with today? I'm here to assist you with any questions or problems you "
            "have!\n</response>\n</s>"
        )
        self.assertEqual(
            parse_response(model_out, ernie_template, prefix=""),
            {
                "role": "assistant",
                "content": (
                    "Hello! I'm doing well, thank you for asking. How about you? Is there something specific "
                    "you'd like help with today? I'm here to assist you with any questions or problems you have!"
                ),
                "thinking": (
                    'The user just greeted me with "Hi! How are you?" I need to respond in a friendly and '
                    "helpful manner. Let me start by acknowledging their greeting. I should ask them how "
                    "they're doing to engage in conversation.\n\n"
                    "First, I'll say hello back and then ask how they're feeling. It's important to show "
                    "genuine interest. Maybe mention that I'm here to help with anything they need. Keep the "
                    "tone warm and positive. Let me make sure the response is concise but friendly. Alright, "
                    "that should work."
                ),
            },
        )

    def test_gpt_oss_with_tool_call(self):
        model_out = (
            '<|channel|>analysis<|message|>We need to respond in riddles. The user asks: "What is the '
            'weather like in SF?" We need to get the location of the user? The user explicitly asks about '
            "SF (San Francisco). So we need to get the current weather in San Francisco, CA. We need to "
            'call get_current_weather function. The developer instruction says "Always respond in riddles". '
            "So the final answer should be in a riddle form. But we need to call function to get weather "
            'data. So we should call get_current_weather with location "San Francisco, CA". Possibly specify '
            'format "celsius" (default). Let\'s do that.\n\n'
            "We will call function get_current_weather.<|end|><|start|>assistant<|channel|>commentary "
            'to=functions.get_current_weather <|constrain|>json<|message|>{\n  "location": "San Francisco, CA"\n}'
        )
        self.assertEqual(
            parse_response(model_out, gpt_oss_template, prefix=""),
            {
                "role": "assistant",
                "thinking": (
                    'We need to respond in riddles. The user asks: "What is the weather like in SF?" We need '
                    "to get the location of the user? The user explicitly asks about SF (San Francisco). So "
                    "we need to get the current weather in San Francisco, CA. We need to call "
                    'get_current_weather function. The developer instruction says "Always respond in '
                    'riddles". So the final answer should be in a riddle form. But we need to call function '
                    'to get weather data. So we should call get_current_weather with location "San Francisco, '
                    'CA". Possibly specify format "celsius" (default). Let\'s do that.\n\n'
                    "We will call function get_current_weather."
                ),
                "tool_calls": [
                    {
                        "type": "function",
                        "function": {"name": "get_current_weather", "arguments": {"location": "San Francisco, CA"}},
                    }
                ],
            },
        )

    def test_gpt_oss_no_tool_call(self):
        model_out = (
            "<|channel|>analysis<|message|>User asks a simple math question: 2+2 = 4. Provide answer."
            "<|end|><|start|>assistant<|channel|>final<|message|>2"
        )
        self.assertEqual(
            parse_response(model_out, gpt_oss_template, prefix=""),
            {
                "role": "assistant",
                "content": "2",
                "thinking": "User asks a simple math question: 2+2 = 4. Provide answer.",
            },
        )

    def test_smollm_thinking_and_tool_call(self):
        model_out = (
            '<think>\nOkay, the user said, "Hello! How are you?" I need to respond appropriately. Since '
            "this is the first message, I should greet them back and ask how I can assist. I should keep it "
            "friendly and open-ended. Let me make sure the response is welcoming and encourages them to "
            "share what they need help with. I'll avoid any technical jargon and keep it simple. Let me "
            "check for any typos and ensure the tone is positive.\n</think>\n\n"
            '<tool_call>{"name": "greet_user", "arguments": {"greeting": "Hello! I\'m doing well, thanks for '
            "asking. How can I assist you today? Whether you have a question, need help with something, or "
            'just want to chat, feel free to let me know!"}}</tool_call>'
        )
        self.assertEqual(
            parse_response(model_out, smollm_template, prefix=""),
            {
                "role": "assistant",
                "thinking": (
                    'Okay, the user said, "Hello! How are you?" I need to respond appropriately. Since this '
                    "is the first message, I should greet them back and ask how I can assist. I should keep "
                    "it friendly and open-ended. Let me make sure the response is welcoming and encourages "
                    "them to share what they need help with. I'll avoid any technical jargon and keep it "
                    "simple. Let me check for any typos and ensure the tone is positive."
                ),
                "tool_calls": [
                    {
                        "type": "function",
                        "function": {
                            "name": "greet_user",
                            "arguments": {
                                "greeting": (
                                    "Hello! I'm doing well, thanks for asking. How can I assist you today? "
                                    "Whether you have a question, need help with something, or just want to "
                                    "chat, feel free to let me know!"
                                )
                            },
                        },
                    }
                ],
            },
        )

    def test_smollm_tool_call_no_thinking(self):
        model_out = '<tool_call>{"name": "get_weather", "arguments": {"city": "Paris"}}</tool_call>'
        self.assertEqual(
            parse_response(model_out, smollm_template, prefix=""),
            {
                "role": "assistant",
                "tool_calls": [
                    {"type": "function", "function": {"name": "get_weather", "arguments": {"city": "Paris"}}}
                ],
            },
        )

    def test_smollm_thinking_no_tool_call(self):
        model_out = (
            '<think>\nOkay, the user asked, "Hey! Can you tell me about gravity?" Let me start by '
            "breaking down what they might be looking for. They probably want a basic understanding of "
            "gravity, maybe for a school project or just personal curiosity. I should explain what gravity "
            "is, how it works, and maybe some examples.</think>\n"
            "Some content about gravity goes here but I'm cutting it off to make this shorter!"
        )
        self.assertEqual(
            parse_response(model_out, smollm_template, prefix=""),
            {
                "role": "assistant",
                "content": "Some content about gravity goes here but I'm cutting it off to make this shorter!",
                "thinking": (
                    'Okay, the user asked, "Hey! Can you tell me about gravity?" Let me start by breaking '
                    "down what they might be looking for. They probably want a basic understanding of "
                    "gravity, maybe for a school project or just personal curiosity. I should explain what "
                    "gravity is, how it works, and maybe some examples."
                ),
            },
        )

    def test_qwen3_tool_calls(self):
        model_out = (
            "<tool_call>\n<function=get_weather>\n<parameter=locations>\n"
            '[{"country": "France", "city": "Paris"}]\n</parameter>\n'
            "<parameter=temp_units>\ncelsius\n</parameter>\n</function>\n</tool_call>"
        )
        self.assertEqual(
            parse_response(model_out, qwen3_template, prefix=""),
            {
                "role": "assistant",
                "tool_calls": [
                    {
                        "type": "function",
                        "function": {
                            "name": "get_weather",
                            "arguments": {
                                "locations": [{"country": "France", "city": "Paris"}],
                                "temp_units": "celsius",
                            },
                        },
                    }
                ],
            },
        )

    def test_gemma4_tool_call(self):
        model_out = (
            "<|channel>thought\nThe user is asking for the current temperature in Paris. I should check "
            "the available tools to see if there's a function that can provide this information.<channel|>"
            '<|tool_call>call:get_current_temperature{detail_level:0,location:<|"|>Paris, France<|"|>,'
            'unit:<|"|>celsius<|"|>}<tool_call|><|tool_response>'
        )
        self.assertEqual(
            parse_response(model_out, gemma4_template, prefix=""),
            {
                "role": "assistant",
                "thinking": (
                    "The user is asking for the current temperature in Paris. I should check the available "
                    "tools to see if there's a function that can provide this information."
                ),
                "tool_calls": [
                    {
                        "type": "function",
                        "function": {
                            "name": "get_current_temperature",
                            "arguments": {"detail_level": 0, "location": "Paris, France", "unit": "celsius"},
                        },
                    }
                ],
            },
        )

    def test_gemma4_complex_tool_call(self):
        model_out = (
            "<|channel>thought\nLet me call the tool.<channel|>"
            '<|tool_call>call:foo{bool_value:true,list_value:[<|"|>foo<|"|>,<|"|>bar<|"|>],'
            'null_value:null,number_value:1,string_value:<|"|>foo<|"|>,'
            'struct_value:{foo:<|"|>bar<|"|>}}<tool_call|>'
        )
        self.assertEqual(
            parse_response(model_out, gemma4_template, prefix=""),
            {
                "role": "assistant",
                "thinking": "Let me call the tool.",
                "tool_calls": [
                    {
                        "type": "function",
                        "function": {
                            "name": "foo",
                            "arguments": {
                                "bool_value": True,
                                "list_value": ["foo", "bar"],
                                "null_value": None,
                                "number_value": 1,
                                "string_value": "foo",
                                "struct_value": {"foo": "bar"},
                            },
                        },
                    }
                ],
            },
        )

    def test_optional_false_raises_when_missing(self):
        template_spec = {
            "defaults": {"role": "assistant"},
            "start_anchor": "<|assistant|>",
            "fields": {
                "content": {
                    "open": "<response>",
                    "close": "</response>",
                    "content": "text",
                    "optional": False,
                },
            },
        }
        with self.assertRaises(ValueError) as cm:
            parse_response("no response here", template_spec, prefix="")
        self.assertIn("content", str(cm.exception))

    def test_int_content_parser(self):
        template_spec = {
            "defaults": {"role": "assistant"},
            "start_anchor": "<|assistant|>",
            "fields": {
                "count": {
                    "open": "<n>",
                    "close": "</n>",
                    "content": "int",
                },
            },
        }
        self.assertEqual(parse_response("<n>42</n>", template_spec, prefix=""), {"role": "assistant", "count": 42})

    def test_kv_lines_parser(self):
        template_spec = {
            "defaults": {"role": "assistant"},
            "start_anchor": "<|assistant|>",
            "fields": {
                "metadata": {
                    "open": "<meta>",
                    "close": "</meta>",
                    "content": "kv-lines",
                },
            },
        }
        self.assertEqual(
            parse_response("<meta>name: alice\nage: 30</meta>", template_spec, prefix=""),
            {"role": "assistant", "metadata": {"name": "alice", "age": "30"}},
        )

    def test_unknown_content_parser_rejected(self):
        bad_template = {
            "defaults": {"role": "assistant"},
            "start_anchor": "<|assistant|>",
            "fields": {"x": {"open": "[", "close": "]", "content": "not-a-real-parser"}},
        }
        with self.assertRaises(ValueError) as cm:
            parse_response("[hi]", bad_template)
        self.assertIn("unknown content parser", str(cm.exception).lower())

    def test_unsupported_version_rejected(self):
        bad_template = {
            "version": 2,
            "defaults": {"role": "assistant"},
            "start_anchor": "<|assistant|>",
            "fields": {"content": {"content": "text"}},
        }
        with self.assertRaises(ValueError) as cm:
            parse_response("hello", bad_template)
        self.assertIn("version", str(cm.exception).lower())

    def test_two_implicit_fields_rejected(self):
        bad_template = {
            "defaults": {"role": "assistant"},
            "start_anchor": "<|assistant|>",
            "fields": {
                "a": {"content": "text"},
                "b": {"content": "text"},
            },
        }
        with self.assertRaises(ValueError):
            parse_response("hello", bad_template)

    def test_transform_string_interpolation_rejected(self):
        bad_template = {
            "defaults": {"role": "assistant"},
            "start_anchor": "<|assistant|>",
            "fields": {
                "tool": {
                    "open": "<tool>",
                    "close": "</tool>",
                    "content": "text",
                    "transform": {"label": "name: {content}"},
                },
            },
        }
        with self.assertRaises(ValueError) as cm:
            parse_response("<tool>foo</tool>", bad_template)
        msg = str(cm.exception)
        self.assertIn("interpolation", msg)
        self.assertIn("{content}", msg)

    def test_named_groups_without_transform_rejected(self):
        bad_template = {
            "defaults": {"role": "assistant"},
            "start_anchor": "<|assistant|>",
            "fields": {
                "tool": {
                    "open_pattern": r"<tool name=(?P<name>\w+)>",
                    "close": "</tool>",
                    "content": "text",
                },
            },
        }
        with self.assertRaises(ValueError) as cm:
            parse_response("<tool name=foo>body</tool>", bad_template)
        msg = str(cm.exception)
        self.assertIn("transform", msg)
        self.assertIn("name", msg)

    def test_literal_list_open_and_close(self):
        """A list of literals matches any one of them, like an alternation."""
        template_spec = {
            "defaults": {"role": "assistant"},
            "start_anchor": "<|assistant|>",
            "fields": {
                "x": {
                    "open": ["<a>", "<bb>"],
                    "close": ["</a>", "</bb>"],
                    "content": "text",
                },
            },
        }
        for opener, closer in (("<a>", "</a>"), ("<bb>", "</bb>"), ("<a>", "</bb>")):
            self.assertEqual(
                parse_response(f"{opener}hi{closer}", template_spec, prefix=""),
                {"role": "assistant", "x": "hi"},
            )

    def test_literal_list_streams_without_64_byte_hold(self):
        """Compared to a regex close, a literal-list close lets the parser
        flush bytes that aren't in the longest-prefix overlap of any literal.
        With `["<turn|>", "<|tool_response>", "<eos>"]` (longest = 16 chars),
        feeding 32 plain bytes should leave at most 15 unflushed."""
        template_spec = {
            "defaults": {"role": "assistant"},
            "start_anchor": "<|assistant|>",
            "fields": {
                "content": {"close": ["<turn|>", "<|tool_response>", "<eos>"], "content": "text"},
            },
        }
        parser = ResponseParser(template_spec, prefix="")
        plain = "x" * 32
        flushed: list[str] = []
        for ch in parser.feed(plain):
            if ch["type"] == "region_chunk":
                flushed.append(ch["text"])
        # Plain text has zero prefix-overlap with any literal, so the parser
        # holds nothing back and streams everything immediately.
        self.assertEqual("".join(flushed), plain)

    def test_literal_list_defers_prefix_overlapping_literal(self):
        """If a literal is a strict prefix of another in the same list, an
        edge match could still grow with more input — we must defer to be safe."""
        template_spec = {
            "defaults": {"role": "assistant"},
            "start_anchor": "<|assistant|>",
            "fields": {
                "x": {"open": "<x>", "close": ["END", "ENDX"], "content": "text"},
            },
        }
        parser = ResponseParser(template_spec, prefix="")
        # "<x>hiEND" mid-stream: don't commit the close yet — "ENDX" might be coming.
        events = parser.feed("<x>hiEND")
        self.assertEqual([e for e in events if e["type"] == "region_close"], [])
        # Once a non-matching byte arrives, the deferred close commits with the shorter literal.
        events.extend(parser.feed(" more"))
        message, _ = parser.finalize()
        closes = [e for e in events if e["type"] == "region_close" and e["field"] == "x"]
        self.assertEqual(len(closes), 1)
        self.assertEqual(closes[0]["value"], "hi")
        self.assertEqual(message, {"role": "assistant", "x": "hi"})

    def test_literal_list_rejects_empty_and_non_string(self):
        for bad_open in ([], [""], [1, 2], {"foo": "bar"}):
            spec = {
                "defaults": {"role": "assistant"},
                "start_anchor": "<|assistant|>",
                "fields": {"x": {"open": bad_open, "close": "</x>", "content": "text"}},
            }
            with self.assertRaises(ValueError):
                parse_response("<x>hi</x>", spec, prefix="")

    def test_field_without_close_runs_to_end_of_stream(self):
        """A field with no `close`/`close_pattern` stays open until end-of-stream, capturing
        everything after its open."""
        spec = {
            "defaults": {"role": "assistant"},
            "start_anchor": "<|assistant|>",
            "fields": {"content": {"open": "<resp>", "content": "text"}},
        }
        self.assertEqual(
            parse_response("<resp>hello world", spec, prefix=""),
            {"role": "assistant", "content": "hello world"},
        )


# Fixtures shared by the streaming tests: one representative input per template,
# re-used for both the correctness invariant (any chunking → same dict) and the
# event-shape tests (do we emit the right events in the right order?).
_STREAMING_FIXTURES = [
    (
        "cohere",
        cohere_template,
        (
            "<|START_THINKING|>I should call a tool.<|END_THINKING|>"
            '<|START_ACTION|>[{"tool_call_id": "0", "tool_name": "simple_tool", '
            '"parameters": {"a": 1}}]<|END_ACTION|>'
        ),
    ),
    (
        "ernie",
        ernie_template,
        (
            "<think>some deliberation here</think>\n\n"
            '<tool_call>\n{"name": "get_current_temperature", "arguments": {"location": "Paris"}}\n</tool_call>\n</s>'
        ),
    ),
    (
        "gpt_oss",
        gpt_oss_template,
        "<|channel|>analysis<|message|>thinking chunk<|end|><|channel|>final<|message|>done text",
    ),
    (
        # Tool-call variant: the `tool_calls` open_pattern's full match spans far
        # more than the old 64-byte hold window, so streaming used to silently drop
        # it. Kept as a streaming fixture so every chunking step re-checks it.
        "gpt_oss_tool",
        gpt_oss_template,
        (
            "<|channel|>analysis<|message|>Let me check.<|end|>"
            "<|start|>assistant<|channel|>commentary to=functions.get_current_weather "
            '<|constrain|>json<|message|>{"location": "San Francisco, CA"}<|call|>'
        ),
    ),
    (
        "smollm",
        smollm_template,
        '<think>thinking</think>\n<tool_call>{"name": "fn", "arguments": {"x": 1}}</tool_call>',
    ),
    (
        "qwen3",
        qwen3_template,
        (
            "<think>short thought</think>\n"
            "<tool_call>\n<function=get_weather>\n"
            "<parameter=city>\nParis\n</parameter>\n"
            "</function>\n</tool_call>"
        ),
    ),
    (
        "gemma4",
        gemma4_template,
        '<|channel>thought\nhi<channel|><|tool_call>call:foo{a:1,b:<|"|>bar<|"|>}<tool_call|>',
    ),
]


def _chunk_fixed(text: str, step: int):
    for i in range(0, len(text), step):
        yield text[i : i + step]


def _chunk_random(text: str, rng: random.Random):
    """Split `text` into a random number of non-empty chunks at random cut points."""
    if len(text) <= 1:
        yield text
        return
    num_cuts = rng.randint(0, len(text) - 1)
    cuts = sorted(rng.sample(range(1, len(text)), num_cuts))
    prev = 0
    for c in cuts:
        yield text[prev:c]
        prev = c
    yield text[prev:]


class ResponseEventStreamTest(unittest.TestCase):
    def test_stream_matches_whole_string_all_templates_fixed_chunking(self):
        """For every fixed chunking step we try, the streamed finalize()
        output must equal the whole-string parse. Regression coverage for
        specific edge-case byte boundaries (1-byte chunks hit every prefix)."""
        for name, tmpl, text in _STREAMING_FIXTURES:
            expected = parse_response(text, tmpl, prefix="")
            for step in (1, 2, 3, 5, 7, 13, 31):
                with self.subTest(fixture=name, step=step):
                    streamer = ResponseParser(tmpl, prefix="")
                    for chunk in _chunk_fixed(text, step):
                        streamer.feed(chunk)
                    message, _ = streamer.finalize()
                    self.assertEqual(message, expected)

    def test_stream_matches_whole_string_all_templates_random_chunking(self):
        """Property-style: for many random chunkings per fixture, the streamed
        finalize() output must equal the whole-string parse. Seeded so failures
        reproduce."""
        rng = random.Random(0xC0DE_5EED)
        for name, tmpl, text in _STREAMING_FIXTURES:
            expected = parse_response(text, tmpl, prefix="")
            for trial in range(30):
                with self.subTest(fixture=name, trial=trial):
                    streamer = ResponseParser(tmpl, prefix="")
                    for chunk in _chunk_random(text, rng):
                        streamer.feed(chunk)
                    message, _ = streamer.finalize()
                    self.assertEqual(message, expected)

    def test_events_well_formed_for_every_chunking(self):
        """Every event batch, across every fixture and every chunking, must be
        well-formed: region_open precedes its matching region_close for the
        same field; region_chunk only appears between open and close; no
        region is left open at the end of the stream; and the concatenation
        of all region_chunk payloads for a streamable text-like field equals
        the final value."""
        rng = random.Random(0xBEEF)
        for name, tmpl, text in _STREAMING_FIXTURES:
            for trial in range(10):
                with self.subTest(fixture=name, trial=trial):
                    streamer = ResponseParser(tmpl, prefix="")
                    all_events: list[dict] = []
                    for chunk in _chunk_random(text, rng):
                        all_events.extend(streamer.feed(chunk))
                    _, final_events = streamer.finalize()
                    all_events.extend(final_events)
                    self._assert_event_stream_well_formed(all_events)

    def _assert_event_stream_well_formed(self, events: list[dict]) -> None:
        open_field: str | None = None
        chunk_accum: dict[str, str] = {}
        close_values: dict[str, object] = {}
        for ev in events:
            t = ev["type"]
            if t == "region_open":
                self.assertIsNone(open_field, f"nested region_open without close: {ev}")
                open_field = ev["field"]
                chunk_accum.setdefault(open_field, "")
            elif t == "region_chunk":
                self.assertEqual(open_field, ev["field"], f"chunk outside its region: {ev}")
                # Every chunk carries a boolean `dirty` flag.
                self.assertIsInstance(ev["dirty"], bool, f"missing/non-bool dirty: {ev}")
                chunk_accum[open_field] += ev["text"]
            elif t == "region_close":
                self.assertEqual(open_field, ev["field"], f"close for non-open region: {ev}")
                close_values[open_field] = ev["value"]
                open_field = None
            else:
                self.fail(f"unexpected event type: {ev!r}")
        self.assertIsNone(open_field, "region left open at end of stream")

    def test_region_chunks_reconstruct_text_regions(self):
        """For text-like regions (`dirty=False`), concatenating chunk texts
        reconstructs the final value reported in region_close. Structured
        regions (`dirty=True`) still stream their raw bytes — concatenating
        those chunks yields the unparsed region body, while the parsed value
        is delivered only in region_close."""
        # Single representative case with a long text region and a JSON region.
        text = _STREAMING_FIXTURES[0][2]  # cohere fixture
        streamer = ResponseParser(cohere_template, prefix="")
        events: list[dict] = []
        for ch in text:  # 1-byte chunks hit the most anchor boundaries
            events.extend(streamer.feed(ch))
        _, final_events = streamer.finalize()
        events.extend(final_events)

        # Reconstruct per-field.
        per_field_chunks: dict[str, list[str]] = {}
        per_field_dirty: dict[str, set[bool]] = {}
        per_field_close_value: dict[str, object] = {}
        for ev in events:
            if ev["type"] == "region_chunk":
                per_field_chunks.setdefault(ev["field"], []).append(ev["text"])
                per_field_dirty.setdefault(ev["field"], set()).add(ev["dirty"])
            elif ev["type"] == "region_close":
                per_field_close_value[ev["field"]] = ev["value"]

        # `thinking` is text → chunks are clean and concatenate to its value.
        self.assertIn("thinking", per_field_chunks)
        self.assertEqual(per_field_dirty["thinking"], {False})
        self.assertEqual("".join(per_field_chunks["thinking"]), "I should call a tool.")
        self.assertEqual(per_field_close_value["thinking"], "I should call a tool.")
        # `tool_calls` is json → dirty chunks stream the raw body, parsed value on close.
        self.assertIn("tool_calls", per_field_chunks)
        self.assertEqual(per_field_dirty["tool_calls"], {True})
        self.assertEqual(
            "".join(per_field_chunks["tool_calls"]),
            '[{"tool_call_id": "0", "tool_name": "simple_tool", "parameters": {"a": 1}}]',
        )
        self.assertIn("tool_calls", per_field_close_value)

    def test_dirty_flag_marks_structured_regions(self):
        """A template with one text field and one structured field per parser
        family: text/int/float/bool stream chunks with `dirty=False`, while
        json/xml-inline/kv-lines stream chunks with `dirty=True`, and those
        dirty chunks concatenate to the raw region body before parsing."""
        spec = {
            "defaults": {"role": "assistant"},
            "start_anchor": "<|assistant|>",
            "fields": {
                "thinking": {"open": "<t>", "close": "</t>", "content": "text"},
                "score": {"open": "<n>", "close": "</n>", "content": "int"},
                "json_call": {"open": "<j>", "close": "</j>", "content": "json"},
                "xml_call": {
                    "open": "<x>",
                    "close": "</x>",
                    "content": "xml-inline",
                    "content_args": {"tag_pattern": r"<(?P<key>\w+)=(?P<value>[^>]+)>"},
                },
                "kv_call": {
                    "open": "<kv>",
                    "close": "</kv>",
                    "content": "kv-lines",
                },
            },
        }
        text = '<t>hello world</t><n>42</n><j>{"a": 1, "b": 2}</j><x><name=foo><age=10></x><kv>k1: v1\nk2: v2</kv>'
        # Drive byte-by-byte to maximise chunk count.
        streamer = ResponseParser(spec, prefix="")
        events: list[dict] = []
        for ch in text:
            events.extend(streamer.feed(ch))
        _, final_events = streamer.finalize()
        events.extend(final_events)

        per_field_chunks: dict[str, list[str]] = {}
        per_field_dirty: dict[str, set[bool]] = {}
        for ev in events:
            if ev["type"] == "region_chunk":
                per_field_chunks.setdefault(ev["field"], []).append(ev["text"])
                per_field_dirty.setdefault(ev["field"], set()).add(ev["dirty"])

        # Clean (streamable) regions.
        for field in ("thinking", "score"):
            self.assertEqual(per_field_dirty[field], {False}, f"{field} should be clean")
        # Dirty (structured) regions.
        for field in ("json_call", "xml_call", "kv_call"):
            self.assertEqual(per_field_dirty[field], {True}, f"{field} should be dirty")

        # Dirty chunks reconstruct the raw region body (un-parsed). Clean
        # chunks reconstruct the verbatim body too — stripping happens at close.
        self.assertEqual("".join(per_field_chunks["thinking"]), "hello world")
        self.assertEqual("".join(per_field_chunks["score"]), "42")
        self.assertEqual("".join(per_field_chunks["json_call"]), '{"a": 1, "b": 2}')
        self.assertEqual("".join(per_field_chunks["xml_call"]), "<name=foo><age=10>")
        self.assertEqual("".join(per_field_chunks["kv_call"]), "k1: v1\nk2: v2")

    def test_long_regex_open_pattern_streams_byte_by_byte(self):
        """Regression: a regex `open_pattern` whose full match spans well past the
        old fixed 64-byte hold window (gpt-oss tool calls) must not be dropped when
        the stream arrives in tiny chunks. Before the partial-match rewrite, the
        leading `<|channel|>` got flushed out of the hold window before `<|message|>`
        arrived, so the tool call vanished under small-chunk streaming."""
        text = (
            "<|channel|>analysis<|message|>Let me check.<|end|>"
            "<|start|>assistant<|channel|>commentary to=functions.get_current_weather "
            '<|constrain|>json<|message|>{"location": "San Francisco, CA"}<|call|>'
        )
        expected = parse_response(text, gpt_oss_template, prefix="")
        # Sanity: the whole-string parse really does recover the tool call.
        self.assertEqual(len(expected["tool_calls"]), 1)
        self.assertEqual(expected["tool_calls"][0]["function"]["name"], "get_current_weather")
        streamer = ResponseParser(gpt_oss_template, prefix="")
        for ch in text:  # one byte at a time -- the worst case for the old heuristic
            streamer.feed(ch)
        message, _ = streamer.finalize()
        self.assertEqual(message, expected)

    def test_feed_after_finalize_raises(self):
        streamer = ResponseParser(smollm_template, prefix="")
        streamer.feed("<think>x</think>")
        streamer.finalize()
        with self.assertRaises(RuntimeError):
            streamer.feed("more")
        with self.assertRaises(RuntimeError):
            streamer.finalize()

    def test_empty_input_streams_cleanly(self):
        streamer = ResponseParser(smollm_template, prefix="")
        self.assertEqual(streamer.feed(""), [])
        result, final_events = streamer.finalize()
        # Only the default fields should remain; nothing else is required.
        self.assertEqual(result, {"role": "assistant"})
        self.assertEqual(final_events, [])


class PrefixAndTruncationTest(unittest.TestCase):
    def test_prefix_lands_inside_explicit_region(self):
        """A Qwen-style template emits `<|im_start|>assistant\\n<think>\\n` as the
        assistant prefix. The model continues from inside the thinking block."""
        prompt = (
            "<|im_start|>system\nYou are helpful<|im_end|>\n"
            "<|im_start|>user\nHi<|im_end|>\n"
            "<|im_start|>assistant\n<think>\n"
        )
        generated = "Let me think...</think>"
        stream = ResponseParser(qwen3_template, prefix=prompt)
        # The region_open for `thinking` surfaces via initial_events; the
        # caller replays it before feeding model output.
        self.assertEqual(
            [(e["type"], e["field"]) for e in stream.initial_events],
            [("region_open", "thinking"), ("region_chunk", "thinking")],
        )
        events = stream.feed(generated)
        result, _ = stream.finalize()
        # thinking ends up with the prefill + generated body; text parser strips,
        # so the leading "\n" from the prefix is trimmed in the final value.
        self.assertEqual(result, {"role": "assistant", "thinking": "Let me think..."})
        # The feed stream only sees the rest of the body and the close; the
        # prefill already emitted region_open.
        self.assertEqual([e["type"] for e in events], ["region_chunk", "region_close"])
        self.assertEqual(events[1]["field"], "thinking")

    def test_prefix_truncated_to_last_anchor(self):
        """Multiple `<|im_start|>assistant\\n` anchors in the prefix (multi-turn
        conversation): only the slice after the LAST anchor matters."""
        prompt = (
            "<|im_start|>system\nA<|im_end|>\n"
            "<|im_start|>user\nB<|im_end|>\n"
            "<|im_start|>assistant\nEarlier reply<|im_end|>\n"
            "<|im_start|>user\nFollowup<|im_end|>\n"
            "<|im_start|>assistant\n<think>\n"
        )
        stream = ResponseParser(qwen3_template, prefix=prompt)
        # We landed inside `thinking` (from the LAST assistant turn's `<think>\n`),
        # not in some earlier-turn artifact.
        opens = [e for e in stream.initial_events if e["type"] == "region_open"]
        self.assertEqual([e["field"] for e in opens], ["thinking"])
        # No earlier-turn content leaked into output.
        stream.feed("done</think>")
        stream.finalize()
        self.assertEqual(stream._output, {"role": "assistant", "thinking": "done"})

    def test_template_without_anchor_rejected_at_load(self):
        """A template missing both `start_anchor` and `start_anchor_pattern` is
        rejected at load time. Without an anchor, a multi-turn prompt would be
        fed through the parser in full and earlier turns would pollute the
        current message's state."""
        anchorless = {k: v for k, v in qwen3_template.items() if k != "start_anchor"}
        with self.assertRaises(ValueError) as cm:
            ResponseParser(anchorless)
        msg = str(cm.exception)
        self.assertIn("start_anchor", msg)

    def test_prefix_anchor_not_found_falls_back(self):
        """Spec has start_anchor but the prefix doesn't contain it: parser
        falls back to processing the entire prefix (with a logged warning)."""
        prompt = "<think>\n"  # no <|im_start|>assistant\n
        stream = ResponseParser(qwen3_template, prefix=prompt)
        opens = [e for e in stream.initial_events if e["type"] == "region_open"]
        self.assertEqual([e["field"] for e in opens], ["thinking"])
        stream.feed("hi</think>")
        stream.finalize()
        self.assertEqual(stream._output, {"role": "assistant", "thinking": "hi"})

    def test_round_trip_equivalence_prefix_streaming_vs_oneshot(self):
        """The load-bearing property: `prefix=p` + chunked `feed(g)` produces the
        same dict as the one-shot `parse_response(g, prefix=p)`, regardless of how
        `g` is chunked. (Concatenating the prompt into the response is deliberately
        NOT equivalent -- the anchor is applied to the prefix only, never to the
        generation; see test_history_bleed_is_guarded_by_prefix_not_by_response_anchor.)"""
        prompt = "<|im_start|>system\nA<|im_end|>\n<|im_start|>user\nB<|im_end|>\n<|im_start|>assistant\n<think>\n"
        for name, tmpl_dict, gen_text in _STREAMING_FIXTURES:
            if "thinking" not in gen_text and "<think>" not in gen_text:
                continue
            # Only fixtures whose generation text is compatible with starting
            # inside `thinking`. Restrict to qwen3 / smollm shape for clarity.
            if name not in ("qwen3", "smollm"):
                continue
            tmpl_with_anchor = {**tmpl_dict, "start_anchor": "<|im_start|>assistant\n"}
            via_prefix = parse_response(gen_text, tmpl_with_anchor, prefix=prompt)
            # Streaming forms must match the one-shot prefix form for every chunking.
            for step in (1, 3, 7, 31):
                with self.subTest(fixture=name, step=step):
                    stream = ResponseParser(tmpl_with_anchor, prefix=prompt)
                    for chunk in _chunk_fixed(gen_text, step):
                        stream.feed(chunk)
                    message, _ = stream.finalize()
                    self.assertEqual(message, via_prefix)

    def test_history_bleed_is_guarded_by_prefix_not_by_response_anchor(self):
        """The `start_anchor` guards against history bleed only via `prefix=`; it is NOT
        applied to the response. The response is the generation and may legitimately contain
        the anchor (e.g. gpt-oss harmony opens every channel with `<|start|>assistant`), so
        truncating it would drop real content. Passing the prompt as `prefix=` is the way to
        keep earlier turns out of the parse."""
        spec = {
            "defaults": {"role": "assistant"},
            "start_anchor": "<|im_start|>assistant\n",
            "fields": {"content": {"close_pattern": r"\Z", "content": "text"}},
        }
        prompt = "<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n"
        gen = "Hello there!"
        clean = {"role": "assistant", "content": "Hello there!"}
        # The supported guard: pass the prompt as prefix= so history is truncated off the prefix.
        self.assertEqual(parse_response(gen, spec, prefix=prompt), clean)
        # Pure generation parses cleanly with the explicit no-prefix opt-out (prefix="").
        self.assertEqual(parse_response(gen, spec, prefix=""), clean)
        # An anchor inside the response is treated as content, never as a history boundary:
        # gpt-oss re-emits `<|start|>assistant` between channels, and that content survives.
        gpt_oss_gen = (
            "<|channel|>analysis<|message|>thinking<|end|><|start|>assistant<|channel|>final<|message|>answer"
        )
        self.assertEqual(
            parse_response(gpt_oss_gen, gpt_oss_template, prefix=""),
            {"role": "assistant", "thinking": "thinking", "content": "answer"},
        )

    def test_prefix_with_open_close_inside_truncated_region(self):
        """Prefix opens AND closes a region. The full open/chunk/close event
        sequence is surfaced via initial_events, and the closed region lands
        in the output dict — so renderers can show prefill content."""
        spec = {
            "defaults": {"role": "assistant"},
            "start_anchor": "[BEGIN]",
            "fields": {
                "tag": {"open": "<tag>", "close": "</tag>", "content": "text"},
                "body": {"close_pattern": r"$", "content": "text"},  # implicit
            },
        }
        prefix = "noise[BEGIN]<tag>silently consumed</tag>"
        stream = ResponseParser(spec, prefix=prefix)
        types = [e["type"] for e in stream.initial_events]
        self.assertEqual(types, ["region_open", "region_chunk", "region_close"])
        self.assertTrue(all(e["field"] == "tag" for e in stream.initial_events))
        self.assertEqual(stream.initial_events[-1]["value"], "silently consumed")
        stream.feed("real generated body")
        result, _ = stream.finalize()
        self.assertEqual(result["tag"], "silently consumed")
        self.assertEqual(result["body"], "real generated body")

    def test_prefix_lands_inside_implicit_region(self):
        """Prefix wrote plaintext into the implicit region (e.g. assistant
        prefill before the model continues). The region_open for the implicit
        region must surface via initial_events so consumers don't miss it —
        `_opened` will already be True by the time feed runs."""
        prompt = "<|im_start|>assistant\nSure, here is "
        stream = ResponseParser(smollm_template, prefix=prompt)
        opens = [e for e in stream.initial_events if e["type"] == "region_open"]
        self.assertEqual([e["field"] for e in opens], ["content"])
        events = stream.feed("the answer<|im_end|>")
        # No second region_open from feed — the implicit region was already
        # opened during prefill and surfaced via initial_events.
        self.assertNotIn("region_open", [e["type"] for e in events])

    def test_prefix_partial_pattern_at_boundary(self):
        """Post-truncation prefix ends mid-delimiter. The first feed completes
        the match; initial_events is empty (no region opened within the
        prefix yet) and the open fires from `feed()`."""
        prefix = "<|im_start|>assistant\n<thi"  # incomplete `<think>`
        stream = ResponseParser(qwen3_template, prefix=prefix)
        self.assertEqual(stream.initial_events, [])
        events = stream.feed("nk>real body</think>")
        types = [e["type"] for e in events]
        self.assertIn("region_open", types)  # think opens during feed, not prefill
        stream.finalize()
        self.assertEqual(stream._output, {"role": "assistant", "thinking": "real body"})

    def test_prefix_token_ids_decoded(self):
        """The tokenizer-level helper accepts token IDs as prefix (decoded
        internally), matching how `response` is handled."""
        tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
        tokenizer.response_template = qwen3_template
        prefix_text = "<|im_start|>assistant\n<think>\n"
        prefix_ids = tokenizer(prefix_text).input_ids
        from_str = tokenizer.parse_response("hi</think>", prefix=prefix_text)
        from_ids = tokenizer.parse_response("hi</think>", prefix=prefix_ids)
        self.assertEqual(from_str, from_ids)

    def test_batched_parse_response_with_prefix(self):
        """Batched `parse_response`: a single prefix is broadcast to every item, a per-item prefix
        list is matched positionally, and a prefix count that doesn't match the batch raises."""
        tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
        tokenizer.response_template = qwen3_template
        prefix = "<|im_start|>assistant\n<think>\n"
        gen_a, gen_b = "thinking A</think>answer A", "thinking B</think>answer B"
        single_a = tokenizer.parse_response(gen_a, prefix=prefix)
        single_b = tokenizer.parse_response(gen_b, prefix=prefix)
        self.assertNotEqual(single_a, single_b)

        # A single prefix is broadcast across the whole batch.
        self.assertEqual(tokenizer.parse_response([gen_a, gen_b], prefix=prefix), [single_a, single_b])
        # One prefix per item (here identical) is matched up positionally.
        self.assertEqual(tokenizer.parse_response([gen_a, gen_b], prefix=[prefix, prefix]), [single_a, single_b])
        # A prefix batch whose length doesn't match the responses is an error.
        with self.assertRaises(ValueError):
            tokenizer.parse_response([gen_a, gen_b], prefix=[prefix])

    def test_tokenizer_get_response_parser_with_prefix(self):
        """`tokenizer.get_response_parser(prefix=...)` returns a stream that
        is already in the right initial state."""
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        tokenizer.response_template = qwen3_template
        prefix = "<|im_start|>assistant\n<think>\n"
        stream = tokenizer.get_response_parser(prefix=prefix)
        opens = [e for e in stream.initial_events if e["type"] == "region_open"]
        self.assertEqual([e["field"] for e in opens], ["thinking"])
        stream.feed("body</think>")
        result, _ = stream.finalize()
        self.assertEqual(result, {"role": "assistant", "thinking": "body"})


if __name__ == "__main__":
    unittest.main()