sglang/test/registered/unit/constrained/test_base_grammar_backend.py

"""
Unit tests for sglang.srt.constrained.base_grammar_backend.

Test Coverage:
- GrammarStats: default values, mutable default isolation
- BaseGrammarObject: default behavior
- InvalidGrammarObject: error message
- BaseGrammarBackend: caching, dispatch routing, unsupported fallback,
  thread pool execution, cache hit/miss
- create_grammar_backend: factory routing, "none" backend, invalid name,
  custom registry, reasoner wrapping
- register_grammar_backend: registration and lookup

Usage:
    python -m pytest test_base_grammar_backend.py -v
"""

import json
import unittest
from concurrent.futures import Future
from unittest.mock import MagicMock, patch

from sglang.srt.arg_groups.overrides import resolution_result
from sglang.srt.constrained.base_grammar_backend import (
    GRAMMAR_BACKEND_REGISTRY,
    BaseGrammarBackend,
    BaseGrammarObject,
    GrammarStats,
    InvalidGrammarObject,
    create_grammar_backend,
    register_grammar_backend,
)
from sglang.srt.runtime_context import get_context  # noqa: E402
from sglang.test.ci.ci_register import register_cpu_ci

register_cpu_ci(2.0, "base-a-test-cpu")


class TestGrammarStats(unittest.TestCase):
    """Test GrammarStats dataclass."""

    def test_defaults(self):
        stats = GrammarStats()
        self.assertIsNone(stats.compilation_time)
        self.assertIsNone(stats.schema_count)
        self.assertIsNone(stats.ebnf_size)
        self.assertFalse(stats.is_cache_hit)
        self.assertFalse(stats.is_grammar_aborted)
        self.assertEqual(stats.tree_traversal_time, [])
        self.assertIsNone(stats.dispatch_type)
        self.assertEqual(stats.num_timeout, 0)

    def test_tree_traversal_time_mutable_default(self):
        """Ensure each instance gets its own list."""
        s1 = GrammarStats()
        s2 = GrammarStats()
        s1.tree_traversal_time.append(0.1)
        self.assertEqual(len(s2.tree_traversal_time), 0)


class TestBaseGrammarObject(unittest.TestCase):
    """Test BaseGrammarObject base class."""


class TestInvalidGrammarObject(unittest.TestCase):
    """Test InvalidGrammarObject."""

    def test_default_error_message(self):
        obj = InvalidGrammarObject()
        self.assertEqual(obj.error_message, "Unknown grammar error")

    def test_custom_error_message(self):
        obj = InvalidGrammarObject("Regex compilation failed")
        self.assertEqual(obj.error_message, "Regex compilation failed")


class TestBaseGrammarBackend(unittest.TestCase):
    """Test BaseGrammarBackend caching and dispatch."""

    def setUp(self):
        self.backend = BaseGrammarBackend()

    def tearDown(self):
        self.backend.executor.shutdown(wait=True)

    def test_cache_hit_returns_copy(self):
        """Cache hit should return a copy of the cached object."""
        mock_copy = BaseGrammarObject()
        obj = MagicMock(spec=BaseGrammarObject)
        obj.copy.return_value = mock_copy

        key = ("json", "schema")
        self.backend.set_cache(key, obj)
        result, cache_hit = self.backend.get_cached_or_future_value(key, False)

        self.assertTrue(cache_hit)
        obj.copy.assert_called_once()
        self.assertIs(result, mock_copy)

    def test_cache_hit_inits_reasoning(self):
        obj = MagicMock(spec=BaseGrammarObject)
        copied = MagicMock(spec=BaseGrammarObject)
        obj.copy.return_value = copied

        key = ("json", "schema")
        self.backend.set_cache(key, obj)
        self.backend.get_cached_or_future_value(key, True)
        copied.maybe_init_reasoning.assert_called_once_with(True)

    def test_cache_miss_returns_future(self):
        key = ("json", "schema")
        result, cache_hit = self.backend.get_cached_or_future_value(key, False)
        self.assertFalse(cache_hit)
        self.assertIsInstance(result, Future)
        # The future should complete (dispatch_json returns InvalidGrammarObject)
        value = result.result(timeout=5)
        self.assertIsInstance(value, InvalidGrammarObject)

    def test_dispatch_fallback_raises(self):
        with self.assertRaises(ValueError):
            self.backend.dispatch_fallback("unknown", "value")

    def test_init_value_dispatch_routes_all_types(self):
        """_init_value_dispatch routes all grammar types to their dispatch methods."""
        cases = [
            ("json", "schema"),
            ("regex", "[a-z]+"),
            ("ebnf", "root ::= 'x'"),
            ("structural_tag", "{}"),
        ]
        for grammar_type, value in cases:
            with self.subTest(grammar_type=grammar_type):
                result = self.backend._init_value_dispatch((grammar_type, value), False)
                self.assertIsInstance(result, InvalidGrammarObject)

    def test_init_value_dispatch_unknown_type_raises(self):
        with self.assertRaises(ValueError):
            self.backend._init_value_dispatch(("unknown_type", "value"), False)

    def test_init_value_dispatch_sets_compilation_time(self):
        """When grammar has stats, compilation_time should be set."""
        mock_grammar = MagicMock(spec=BaseGrammarObject)
        mock_grammar.grammar_stats = GrammarStats()
        self.backend.dispatch_json = MagicMock(return_value=mock_grammar)

        result = self.backend._init_value_dispatch(("json", "schema"), False)
        self.assertIsNotNone(result.grammar_stats.compilation_time)
        self.assertGreater(result.grammar_stats.compilation_time, 0)

    def test_init_value_dispatch_no_stats(self):
        """When grammar has no stats, should not crash."""
        mock_grammar = MagicMock(spec=BaseGrammarObject)
        mock_grammar.grammar_stats = None
        self.backend.dispatch_json = MagicMock(return_value=mock_grammar)
        # Should not raise
        self.backend._init_value_dispatch(("json", "schema"), False)

    def test_reset_then_miss(self):
        """After reset, previously cached keys should be misses."""
        key = ("json", "schema")
        obj = MagicMock(spec=BaseGrammarObject)
        obj.copy.return_value = obj
        self.backend.set_cache(key, obj)

        _, hit = self.backend.get_cached_or_future_value(key, False)
        self.assertTrue(hit)

        self.backend.reset()
        result, hit = self.backend.get_cached_or_future_value(key, False)
        self.assertFalse(hit)
        self.assertIsInstance(result, Future)

    def test_dispatch_fallback_error_message_content(self):
        """dispatch_fallback error should include the key type and value."""
        with self.assertRaises(ValueError) as ctx:
            self.backend.dispatch_fallback("custom_type", "custom_value")
        self.assertIn("custom_type", str(ctx.exception))
        self.assertIn("custom_value", str(ctx.exception))

    def test_init_value_dispatch_none_grammar(self):
        """When dispatch returns None, should not crash on stats check."""
        self.backend.dispatch_json = MagicMock(return_value=None)
        result = self.backend._init_value_dispatch(("json", "schema"), False)
        self.assertIsNone(result)

    def test_cache_miss_duplicate_key_submits_separate_futures(self):
        """Two cache misses for the same key each get their own Future.

        The backend does not deduplicate in-flight compilations — that is
        handled at the GrammarManager level via grammar_queue. Each call
        to get_cached_or_future_value with an uncached key submits a new
        task to the executor."""
        key = ("json", "schema")
        result1, hit1 = self.backend.get_cached_or_future_value(key, False)
        result2, hit2 = self.backend.get_cached_or_future_value(key, False)

        self.assertFalse(hit1)
        self.assertFalse(hit2)
        self.assertIsInstance(result1, Future)
        self.assertIsInstance(result2, Future)
        # They are independent futures, not shared
        self.assertIsNot(result1, result2)

        # Both should complete successfully
        self.assertIsInstance(result1.result(timeout=5), InvalidGrammarObject)
        self.assertIsInstance(result2.result(timeout=5), InvalidGrammarObject)


class TestRegisterGrammarBackend(unittest.TestCase):
    """Test grammar backend registry."""

    def setUp(self):
        self._saved = dict(GRAMMAR_BACKEND_REGISTRY)

    def tearDown(self):
        GRAMMAR_BACKEND_REGISTRY.clear()
        GRAMMAR_BACKEND_REGISTRY.update(self._saved)

    def test_overwrite_registration(self):
        register_grammar_backend("dup", lambda *a: "first")
        register_grammar_backend("dup", lambda *a: "second")
        self.assertEqual(
            GRAMMAR_BACKEND_REGISTRY["dup"](None, None, None, None), "second"
        )


class TestCreateGrammarBackend(unittest.TestCase):
    """Test create_grammar_backend factory function."""

    def setUp(self):
        self._saved = dict(GRAMMAR_BACKEND_REGISTRY)

    def tearDown(self):
        GRAMMAR_BACKEND_REGISTRY.clear()
        GRAMMAR_BACKEND_REGISTRY.update(self._saved)

    def _publish(self, **fields):
        """Set the config the factory reads.

        The factory takes every config value off the published bags, so a test
        that sets one on the handed object would be setting something the
        factory does not read -- which is how a mismatch between the two used
        to stay invisible here.
        """
        override = get_context().override_server_args(**fields)
        override.install()
        self.addCleanup(override.restore)

    def _make_server_args(
        self,
        backend="none",
        reasoning_parser=None,
        enable_strict_thinking=False,
        **fields,
    ):
        published = {
            "grammar_backend": backend,
            "reasoning_parser": reasoning_parser,
            "enable_strict_thinking": enable_strict_thinking,
            "constrained_json_whitespace_pattern": None,
            "constrained_json_disable_any_whitespace": False,
            "constrained_json_max_whitespace_cnt": None,
        }
        published.update(fields)
        self._publish(**published)
        # Handed on to plugin-registered backends; not a config source.
        args = MagicMock()
        args.override = lambda source, **updates: [
            setattr(args, key, value) for key, value in updates.items()
        ]
        return args

    def test_none_backend_returns_none(self):
        args = self._make_server_args("none")
        result = create_grammar_backend(args, None, 32000)
        self.assertIsNone(result)

    def test_none_backend_with_strict_thinking_raises(self):
        args = self._make_server_args("none", enable_strict_thinking=True)
        with self.assertRaisesRegex(ValueError, "enable-strict-thinking"):
            create_grammar_backend(args, None, 32000)

    def test_invalid_backend_raises(self):
        args = self._make_server_args("nonexistent_backend")
        with self.assertRaises(ValueError):
            create_grammar_backend(args, None, 32000)

    def test_custom_backend_receives_args(self):
        received = {}

        def init_fn(server_args, tokenizer, vocab_size, eos_token_ids):
            received["server_args"] = server_args
            received["tokenizer"] = tokenizer
            received["vocab_size"] = vocab_size
            received["eos_token_ids"] = eos_token_ids
            return MagicMock()

        register_grammar_backend("capture", init_fn)
        args = self._make_server_args("capture")
        create_grammar_backend(args, "my_tok", 50000, {0, 2})
        self.assertEqual(received["tokenizer"], "my_tok")
        self.assertEqual(received["vocab_size"], 50000)
        self.assertEqual(received["eos_token_ids"], {0, 2})

    def test_custom_backend_skips_reasoner_wrapping(self):
        """Custom registered backends return directly, bypassing reasoner wrapping."""
        mock_inner = MagicMock(spec=BaseGrammarBackend)
        register_grammar_backend("inner_r", lambda *a: mock_inner)

        args = self._make_server_args("inner_r", reasoning_parser="deepseek-r1")
        tokenizer = MagicMock()

        result = create_grammar_backend(args, tokenizer, 32000)
        # Custom backends return early, no reasoner wrapping applied
        self.assertIs(result, mock_inner)

    @patch("sglang.srt.constrained.outlines_backend.OutlinesGrammarBackend")
    def test_outlines_backend(self, mock_outlines_cls):
        mock_backend = MagicMock(spec=BaseGrammarBackend)
        mock_outlines_cls.return_value = mock_backend
        args = self._make_server_args(
            "outlines", constrained_json_whitespace_pattern=r"\s*"
        )

        result = create_grammar_backend(args, "tok", 32000)
        mock_outlines_cls.assert_called_once_with("tok", whitespace_pattern=r"\s*")
        self.assertIs(result, mock_backend)

    @patch("sglang.srt.constrained.xgrammar_backend.XGrammarGrammarBackend")
    def test_xgrammar_backend(self, mock_xgrammar_cls):
        mock_backend = MagicMock(spec=BaseGrammarBackend)
        mock_xgrammar_cls.return_value = mock_backend
        args = self._make_server_args(
            "xgrammar", constrained_json_disable_any_whitespace=True
        )

        result = create_grammar_backend(args, "tok", 32000, {1, 2})
        mock_xgrammar_cls.assert_called_once_with(
            "tok",
            vocab_size=32000,
            model_eos_token_ids=[1, 2],
            any_whitespace=False,
            max_whitespace_cnt=None,
        )
        self.assertIs(result, mock_backend)

    @patch("sglang.srt.constrained.xgrammar_backend.XGrammarGrammarBackend")
    def test_xgrammar_unsupported_tokenizer_falls_back_to_none(self, mock_xgrammar_cls):
        from sglang.srt.constrained.xgrammar_backend import TokenizerNotSupportedError
        from sglang.srt.runtime_context import get_context, get_exec

        mock_xgrammar_cls.side_effect = TokenizerNotSupportedError(
            "unsupported tokenizer"
        )
        override = get_context().override_server_args(grammar_backend="xgrammar")
        server_args = override.install()
        self.addCleanup(override.restore)

        self.assertIsNone(create_grammar_backend(server_args, "tok", 32000, {1}))
        self.assertEqual(get_exec().kernel.grammar_backend, "none")
        self.assertEqual(
            get_context().resolved_server_args_dict()["grammar_backend"], "none"
        )
        # The record is not written any more: what the caller asked for is a
        # declaration on it, and the runtime fallback to "none" lives in the bag
        # (asserted above). The two are meant to differ here.
        self.assertEqual(resolution_result(server_args, "grammar_backend"), "xgrammar")

    @patch("sglang.srt.constrained.llguidance_backend.GuidanceBackend")
    def test_llguidance_backend(self, mock_guidance_cls):
        mock_backend = MagicMock(spec=BaseGrammarBackend)
        mock_guidance_cls.return_value = mock_backend
        args = self._make_server_args(
            "llguidance",
            constrained_json_disable_any_whitespace=False,
            constrained_json_whitespace_pattern=r"\s+",
        )

        result = create_grammar_backend(args, "tok", 32000, {1, 2})
        mock_guidance_cls.assert_called_once_with(
            tokenizer="tok",
            any_whitespace=True,
            whitespace_pattern=r"\s+",
            n_vocab=32000,
            eos_token_ids={1, 2},
        )
        self.assertIs(result, mock_backend)

    @patch("sglang.srt.constrained.outlines_backend.OutlinesGrammarBackend")
    def test_reasoner_wrapping_on_builtin_backend(self, mock_outlines_cls):
        """Non-custom backends get wrapped with ReasonerGrammarBackend."""
        from sglang.srt.constrained.reasoner_grammar_backend import (
            ReasonerGrammarBackend,
        )

        mock_backend = MagicMock(spec=BaseGrammarBackend)
        mock_backend.is_support_token_filter = False
        mock_outlines_cls.return_value = mock_backend
        args = self._make_server_args("outlines", reasoning_parser="deepseek-r1")
        tokenizer = MagicMock()
        # encode must return a single-token list for think_start/end tokens
        tokenizer.encode.return_value = [42]

        result = create_grammar_backend(args, tokenizer, 32000, think_end_ids=[42])
        self.assertIsInstance(result, ReasonerGrammarBackend)
        self.assertIs(result.grammar_backend, mock_backend)
        self.assertEqual(result.think_end_ids, [42])

    @patch("sglang.srt.constrained.outlines_backend.OutlinesGrammarBackend")
    def test_no_reasoner_wrapping_without_think_end_ids(self, mock_outlines_cls):
        mock_backend = MagicMock(spec=BaseGrammarBackend)
        mock_outlines_cls.return_value = mock_backend
        args = self._make_server_args("outlines", reasoning_parser="deepseek-r1")
        tokenizer = MagicMock(spec=[])

        result = create_grammar_backend(args, tokenizer, 32000, think_end_ids=None)
        self.assertIs(result, mock_backend)

    @patch("sglang.srt.constrained.outlines_backend.OutlinesGrammarBackend")
    def test_no_reasoner_wrapping_without_reasoning_parser(self, mock_outlines_cls):
        mock_backend = MagicMock(spec=BaseGrammarBackend)
        mock_outlines_cls.return_value = mock_backend
        args = self._make_server_args("outlines", reasoning_parser=None)
        tokenizer = MagicMock()

        result = create_grammar_backend(args, tokenizer, 32000, think_end_ids=[42])
        self.assertIs(result, mock_backend)

    @patch("sglang.srt.constrained.xgrammar_backend.XGrammarGrammarBackend")
    def test_xgrammar_eos_none(self, mock_xgrammar_cls):
        """eos_token_ids=None should pass None, not an empty list."""
        mock_xgrammar_cls.return_value = MagicMock(spec=BaseGrammarBackend)
        args = self._make_server_args("xgrammar")

        create_grammar_backend(args, "tok", 32000, None)
        _, kwargs = mock_xgrammar_cls.call_args
        self.assertIsNone(kwargs["model_eos_token_ids"])


class TestLlguidanceStructuralTagTriggerPairing(unittest.TestCase):
    """Bug regression: dispatch_structural_tag paired EVERY structure with
    triggers[0]. Detectors with per-tool triggers (Inkling emits
    <|message_model|>{name}<|content_invoke_tool_json|> per tool) produce
    multiple distinct triggers, and llguidance's StructTag asserts
    begin.startswith(trigger) — so any multi-tool constrained request
    compiled to InvalidGrammarObject."""

    def test_each_structure_pairs_with_its_own_trigger(self):
        import json

        from sglang.srt.constrained.llguidance_backend import GuidanceBackend

        backend = object.__new__(GuidanceBackend)
        backend._from_serialized = lambda serialized: serialized
        begins = [
            '<|message_model|>alpha<|content_invoke_tool_json|>{"name":"alpha","args":',
            '<|message_model|>beta<|content_invoke_tool_json|>{"name":"beta","args":',
        ]
        key = json.dumps(
            {
                "type": "structural_tag",
                "structures": [
                    {
                        "begin": begin,
                        "schema": {"type": "object"},
                        "end": "<|end_message|>",
                    }
                    for begin in begins
                ],
                "triggers": [
                    "<|message_model|>alpha<|content_invoke_tool_json|>",
                    "<|message_model|>beta<|content_invoke_tool_json|>",
                ],
            }
        )
        result = backend.dispatch_structural_tag(key)
        self.assertNotIsInstance(result, InvalidGrammarObject)


class TestNulByteGrammarRejection(unittest.TestCase):
    def setUp(self):
        self.backend = BaseGrammarBackend()

    def tearDown(self):
        self.backend.executor.shutdown(wait=True)

    def test_nul_payload_never_reaches_backend(self):
        cases = [
            ("regex", "\x00\x01\x02\x1f"),
            ("regex", "\x00"),
            # a non-leading NUL truncates the pattern instead of crashing; pinned
            # so the guard cannot be narrowed to startswith()
            ("regex", "a\x00b"),
            ("ebnf", "root ::= \x00"),
            ("structural_tag", '{"triggers": ["\x00"]}'),
            # escaped forms: no raw NUL byte anywhere in the key string
            ("json", '{"type":"string","pattern":"\\u0000"}'),
            (
                "json",
                '{"type":"object","properties":'
                '{"f":{"type":"string","pattern":"\\u0000x"}}}',
            ),
            ("structural_tag", '{"format":{"pattern":"\\u0000"}}'),
        ]
        for key_type, key_string in cases:
            with self.subTest(key_type=key_type, key_string=repr(key_string)):
                dispatch = MagicMock()
                setattr(self.backend, f"dispatch_{key_type}", dispatch)

                result = self.backend._init_value_dispatch(
                    (key_type, key_string), False
                )

                self.assertIsInstance(result, InvalidGrammarObject)
                dispatch.assert_not_called()

    def test_nul_free_payload_still_dispatches(self):
        cases = [
            ("regex", "[0-9]+"),
            ("regex", r"\x00"),  # escaped in the pattern text, not a raw NUL byte
            ("regex", "\x01\x02\x1f"),  # other control bytes are not the trigger
            ("json", json.dumps({"type": "string", "pattern": "[0-9]+"})),
            # malformed json must reach the backend and get its normal error,
            # not be swallowed by the NUL scan's decode failure
            ("json", "{not valid json"),
        ]
        for key_type, key_string in cases:
            with self.subTest(key_type=key_type, key_string=repr(key_string)):
                dispatch = MagicMock(return_value=BaseGrammarObject())
                setattr(self.backend, f"dispatch_{key_type}", dispatch)

                self.backend._init_value_dispatch((key_type, key_string), False)

                dispatch.assert_called_once_with(key_string)


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