/test/registered/unit/constrained/test_grammar_manager.py

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

Test Coverage:
- GrammarManager initialization, queue management, len, clear
- process_req_with_grammar: dispatch by constraint type (json, regex, ebnf,
  structural_tag), no-constraint requests, no-backend error, cache hits,
  cached invalid grammar abort
- abort_requests: single abort, abort all, future cancellation
- get_ready_grammar_requests: future completion, invalid grammar handling,
  timeout with max poll iterations, aborted request handling, queue cleanup

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

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

from sglang.srt.constrained.base_grammar_backend import (
    BaseGrammarBackend,
    BaseGrammarObject,
    InvalidGrammarObject,
)
from sglang.srt.constrained.grammar_manager import GrammarManager
from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject
from sglang.srt.distributed.communication_tags import P2PTag
from sglang.test.ci.ci_register import register_cpu_ci

register_cpu_ci(2.0, "base-a-test-cpu")
register_cpu_ci(est_time=7, suite="base-c-test-cpu")


def _make_scheduler(grammar_backend_name="none", skip_tokenizer=False):
    """Create a mock scheduler with necessary attributes."""
    scheduler = MagicMock()
    scheduler.server_args.grammar_backend = grammar_backend_name
    scheduler.server_args.skip_tokenizer_init = skip_tokenizer
    scheduler.server_args.reasoning_parser = None
    scheduler.server_args.constrained_json_whitespace_pattern = None
    scheduler.server_args.constrained_json_disable_any_whitespace = False

    # Distributed group mocks
    scheduler.dp_tp_cpu_group = MagicMock()
    scheduler.dp_tp_group.world_size = 1
    scheduler.dp_tp_group.first_rank = 0
    scheduler.dp_tp_group.is_first_rank = True
    scheduler.ps.pp_rank = 0
    scheduler.ps.pp_size = 1
    scheduler.pp_group = None

    return scheduler


def _make_req(
    json_schema=None,
    regex=None,
    ebnf=None,
    structural_tag=None,
    rid="req-1",
    custom_params=None,
):
    """Create a mock request with sampling params."""
    req = MagicMock()
    req.rid = rid
    req.sampling_params.json_schema = json_schema
    req.sampling_params.regex = regex
    req.sampling_params.ebnf = ebnf
    req.sampling_params.structural_tag = structural_tag
    req.sampling_params.custom_params = custom_params
    req.require_reasoning = False
    req.grammar = None
    req.grammar_key = None
    req.grammar_wait_ct = 0
    req.finished.return_value = False
    return req


class TestGrammarManagerInit(unittest.TestCase):
    """Test GrammarManager initialization."""

    @patch("sglang.srt.constrained.grammar_manager.create_grammar_backend")
    def test_init_with_backend(self, mock_create):
        mock_create.return_value = MagicMock(spec=BaseGrammarBackend)
        scheduler = _make_scheduler("xgrammar")
        scheduler.server_args.skip_tokenizer_init = False

        mgr = GrammarManager(scheduler)
        self.assertIsNotNone(mgr.grammar_backend)
        self.assertEqual(len(mgr), 0)

    def test_init_skip_tokenizer(self):
        scheduler = _make_scheduler(skip_tokenizer=True)
        mgr = GrammarManager(scheduler)
        self.assertIsNone(mgr.grammar_backend)

    @patch("sglang.srt.constrained.grammar_manager.create_grammar_backend")
    def test_len_and_has_waiting(self, mock_create):
        mock_create.return_value = None
        scheduler = _make_scheduler()
        mgr = GrammarManager(scheduler)
        self.assertEqual(len(mgr), 0)
        self.assertFalse(mgr.has_waiting_grammars())

    @patch("sglang.srt.constrained.grammar_manager.create_grammar_backend")
    def test_clear_resets_backend(self, mock_create):
        mock_backend = MagicMock(spec=BaseGrammarBackend)
        mock_create.return_value = mock_backend
        scheduler = _make_scheduler()
        scheduler.server_args.skip_tokenizer_init = False

        mgr = GrammarManager(scheduler)
        mgr.clear()
        mock_backend.reset.assert_called_once()

    @patch("sglang.srt.constrained.grammar_manager.create_grammar_backend")
    def test_clear_no_backend(self, mock_create):
        mock_create.return_value = None
        scheduler = _make_scheduler()
        mgr = GrammarManager(scheduler)
        mgr.clear()  # Should not raise


class TestProcessReqWithGrammar(unittest.TestCase):
    """Test process_req_with_grammar dispatch and caching."""

    def _make_mgr(self):
        scheduler = _make_scheduler()
        scheduler.server_args.skip_tokenizer_init = True
        mgr = GrammarManager(scheduler)
        mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend)
        return mgr

    def test_no_constraint_returns_false(self):
        mgr = self._make_mgr()
        req = _make_req()  # No constraints
        result = mgr.process_req_with_grammar(req)
        self.assertFalse(result)
        self.assertEqual(len(mgr.grammar_queue), 0)

    def test_json_schema_cache_miss(self):
        mgr = self._make_mgr()
        future = Future()
        mgr.grammar_backend.get_cached_or_future_value.return_value = (future, False)

        req = _make_req(json_schema='{"type": "object"}')
        result = mgr.process_req_with_grammar(req)

        self.assertTrue(result)
        self.assertEqual(len(mgr.grammar_queue), 1)
        self.assertEqual(req.grammar_key, ("json", '{"type": "object"}'))

    def test_regex_cache_miss(self):
        mgr = self._make_mgr()
        future = Future()
        mgr.grammar_backend.get_cached_or_future_value.return_value = (future, False)

        req = _make_req(regex="[a-z]+")
        result = mgr.process_req_with_grammar(req)

        self.assertTrue(result)
        self.assertEqual(req.grammar_key, ("regex", "[a-z]+"))

    def test_ebnf_cache_miss(self):
        mgr = self._make_mgr()
        future = Future()
        mgr.grammar_backend.get_cached_or_future_value.return_value = (future, False)

        req = _make_req(ebnf="root ::= 'hello'")
        result = mgr.process_req_with_grammar(req)

        self.assertTrue(result)
        self.assertEqual(req.grammar_key, ("ebnf", "root ::= 'hello'"))

    def test_structural_tag_cache_miss(self):
        mgr = self._make_mgr()
        future = Future()
        mgr.grammar_backend.get_cached_or_future_value.return_value = (future, False)

        req = _make_req(structural_tag='{"structures": [], "triggers": []}')
        result = mgr.process_req_with_grammar(req)

        self.assertTrue(result)
        self.assertEqual(
            req.grammar_key,
            ("structural_tag", '{"structures": [], "triggers": []}'),
        )

    def test_cache_hit_returns_false(self):
        """Cache hit should NOT add to grammar queue."""
        mgr = self._make_mgr()
        grammar_obj = MagicMock(spec=BaseGrammarObject)
        mgr.grammar_backend.get_cached_or_future_value.return_value = (
            grammar_obj,
            True,
        )

        req = _make_req(json_schema='{"type": "object"}')
        result = mgr.process_req_with_grammar(req)

        self.assertFalse(result)
        self.assertEqual(len(mgr.grammar_queue), 0)
        self.assertIs(req.grammar, grammar_obj)

    def test_cache_hit_invalid_grammar_aborts(self):
        """Cache hit with InvalidGrammarObject should abort the request."""
        mgr = self._make_mgr()
        invalid = InvalidGrammarObject("bad schema")
        mgr.grammar_backend.get_cached_or_future_value.return_value = (invalid, True)

        req = _make_req(json_schema="bad")
        result = mgr.process_req_with_grammar(req)

        self.assertFalse(result)
        req.set_finish_with_abort.assert_called_once()
        self.assertIn("bad schema", req.set_finish_with_abort.call_args[0][0])

    def test_no_backend_aborts(self):
        """No grammar backend should abort request."""
        scheduler = _make_scheduler()
        scheduler.server_args.skip_tokenizer_init = True
        mgr = GrammarManager(scheduler)
        mgr.grammar_backend = None

        req = _make_req(json_schema='{"type": "object"}')
        result = mgr.process_req_with_grammar(req)

        self.assertFalse(result)
        req.set_finish_with_abort.assert_called_once()
        self.assertIn("not supported", req.set_finish_with_abort.call_args[0][0])

    def test_json_takes_priority_over_other_constraints(self):
        """When json_schema is set, it should be used regardless of other fields."""
        mgr = self._make_mgr()
        future = Future()
        mgr.grammar_backend.get_cached_or_future_value.return_value = (future, False)

        req = _make_req(json_schema='{"type": "object"}', regex="[a-z]+")
        mgr.process_req_with_grammar(req)
        self.assertEqual(req.grammar_key, ("json", '{"type": "object"}'))

    def test_require_reasoning_forwarded_to_backend(self):
        """require_reasoning from the request should be passed to the backend."""
        mgr = self._make_mgr()
        grammar_obj = MagicMock(spec=BaseGrammarObject)
        mgr.grammar_backend.get_cached_or_future_value.return_value = (
            grammar_obj,
            True,
        )

        req = _make_req(json_schema="schema")
        req.require_reasoning = True
        mgr.process_req_with_grammar(req)

        mgr.grammar_backend.get_cached_or_future_value.assert_called_once_with(
            ("json", "schema"), True
        )

    def test_has_waiting_grammars_after_enqueue(self):
        mgr = self._make_mgr()
        future = Future()
        mgr.grammar_backend.get_cached_or_future_value.return_value = (future, False)

        self.assertFalse(mgr.has_waiting_grammars())
        req = _make_req(json_schema="schema")
        mgr.process_req_with_grammar(req)
        self.assertTrue(mgr.has_waiting_grammars())
        self.assertEqual(len(mgr), 1)

    def test_cache_hit_applies_request_thinking_budget(self):
        mgr = self._make_mgr()
        grammar_obj = ReasonerGrammarObject(
            grammar=None, think_end_id=0, max_think_tokens=99
        )
        mgr.grammar_backend.get_cached_or_future_value.return_value = (
            grammar_obj,
            True,
        )

        req = _make_req(
            json_schema="schema",
            custom_params={"thinking_budget": 7},
        )
        mgr.process_req_with_grammar(req)

        self.assertEqual(req.grammar.max_think_tokens, 7)

    def test_strict_reasoning_grammar_applies_request_thinking_budget(self):
        mgr = self._make_mgr()
        mgr._enable_strict_thinking = True
        grammar_obj = ReasonerGrammarObject(
            grammar=None, think_end_id=0, max_think_tokens=99
        )
        mgr.grammar_backend.init_strict_reasoning_grammar.return_value = grammar_obj

        req = _make_req(custom_params={"thinking_budget": 3})
        req.require_reasoning = True
        mgr.process_req_with_grammar(req)

        self.assertIs(req.grammar, grammar_obj)
        self.assertEqual(req.grammar.max_think_tokens, 3)


class TestAbortRequests(unittest.TestCase):
    """Test abort_requests handling."""

    def _make_mgr_with_queue(self):
        scheduler = _make_scheduler()
        scheduler.server_args.skip_tokenizer_init = True
        mgr = GrammarManager(scheduler)
        mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend)
        return mgr

    def test_abort_by_rid_prefix(self):
        mgr = self._make_mgr_with_queue()
        req = _make_req(rid="req-123")
        future = MagicMock(spec=Future)
        req.grammar = future
        mgr.grammar_queue.append(req)

        abort_req = MagicMock()
        abort_req.abort_all = False
        abort_req.rid = "req-123"

        mgr.abort_requests(abort_req)
        future.cancel.assert_called_once()
        req.set_finish_with_abort.assert_called_once()

    def test_abort_non_matching_rid(self):
        mgr = self._make_mgr_with_queue()
        req = _make_req(rid="req-999")
        req.grammar = MagicMock(spec=Future)
        mgr.grammar_queue.append(req)

        abort_req = MagicMock()
        abort_req.abort_all = False
        abort_req.rid = "req-123"

        mgr.abort_requests(abort_req)
        req.set_finish_with_abort.assert_not_called()

    def test_abort_all(self):
        mgr = self._make_mgr_with_queue()
        reqs = []
        for i in range(3):
            req = _make_req(rid=f"req-{i}")
            req.grammar = MagicMock(spec=Future)
            mgr.grammar_queue.append(req)
            reqs.append(req)

        abort_req = MagicMock()
        abort_req.abort_all = True
        abort_req.rid = ""

        mgr.abort_requests(abort_req)
        for req in reqs:
            req.set_finish_with_abort.assert_called_once()

    def test_abort_empty_queue(self):
        """Aborting on an empty queue should not raise."""
        mgr = self._make_mgr_with_queue()
        abort_req = MagicMock()
        abort_req.abort_all = True
        abort_req.rid = ""
        mgr.abort_requests(abort_req)  # Should not raise

    def test_abort_prefix_match(self):
        """rid.startswith means prefix matching, not exact matching."""
        mgr = self._make_mgr_with_queue()
        req = _make_req(rid="req-123-suffix")
        req.grammar = MagicMock(spec=Future)
        mgr.grammar_queue.append(req)

        abort_req = MagicMock()
        abort_req.abort_all = False
        abort_req.rid = "req-123"

        mgr.abort_requests(abort_req)
        req.set_finish_with_abort.assert_called_once()


class TestGetReadyGrammarRequests(unittest.TestCase):
    """Test get_ready_grammar_requests polling and result handling."""

    def _make_mgr(self):
        scheduler = _make_scheduler()
        scheduler.server_args.skip_tokenizer_init = True
        mgr = GrammarManager(scheduler)
        mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend)
        # Use very short poll interval for tests
        mgr.SGLANG_GRAMMAR_POLL_INTERVAL = 0.01
        mgr.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = 3
        return mgr

    def test_ready_future_returns_req(self):
        mgr = self._make_mgr()

        grammar_obj = MagicMock(spec=BaseGrammarObject)
        grammar_obj.copy.return_value = grammar_obj
        future = Future()
        future.set_result(grammar_obj)

        req = _make_req(json_schema="schema")
        req.grammar = future
        req.grammar_key = ("json", "schema")
        mgr.grammar_queue.append(req)

        result = mgr.get_ready_grammar_requests()
        self.assertEqual(len(result), 1)
        self.assertIs(result[0], req)
        self.assertIs(req.grammar, grammar_obj)
        # Cache should be set
        mgr.grammar_backend.set_cache.assert_called_once()
        # Queue should be empty
        self.assertEqual(len(mgr.grammar_queue), 0)

    def test_invalid_grammar_aborts_req(self):
        mgr = self._make_mgr()

        invalid = InvalidGrammarObject("compile error")
        invalid_copy = InvalidGrammarObject("compile error")
        invalid.copy = MagicMock(return_value=invalid_copy)
        future = Future()
        future.set_result(invalid)

        req = _make_req(json_schema="bad")
        req.grammar = future
        req.grammar_key = ("json", "bad")
        mgr.grammar_queue.append(req)

        result = mgr.get_ready_grammar_requests()
        self.assertEqual(len(result), 1)
        req.set_finish_with_abort.assert_called_once()
        self.assertIn("compile error", req.set_finish_with_abort.call_args[0][0])

    def test_aborted_req_removed_from_queue(self):
        mgr = self._make_mgr()

        req = _make_req(json_schema="schema")
        req.finished.return_value = True  # Already aborted
        req.grammar = None
        mgr.grammar_queue.append(req)

        result = mgr.get_ready_grammar_requests()
        self.assertEqual(len(result), 1)
        self.assertEqual(len(mgr.grammar_queue), 0)

    def test_timeout_aborts_req(self):
        mgr = self._make_mgr()
        mgr.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = 1

        future = Future()  # Never completes
        req = _make_req(json_schema="slow")
        req.grammar = future
        req.grammar_key = ("json", "slow")
        req.grammar_wait_ct = 0
        mgr.grammar_queue.append(req)

        # First call: not ready, increments wait_ct to 1 (== max_poll)
        result = mgr.get_ready_grammar_requests()
        # Should timeout and abort
        self.assertEqual(len(result), 1)
        req.set_finish_with_abort.assert_called_once()
        self.assertIn("timed out", req.set_finish_with_abort.call_args[0][0])
        # Cache should store InvalidGrammarObject for timeout
        mgr.grammar_backend.set_cache.assert_called_once()
        cached_key, cached_val = mgr.grammar_backend.set_cache.call_args[0]
        self.assertEqual(cached_key, ("json", "slow"))
        self.assertIsInstance(cached_val, InvalidGrammarObject)

    def test_pending_future_stays_in_queue(self):
        """Futures that aren't done stay in the queue."""
        mgr = self._make_mgr()
        mgr.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = 100  # High to avoid timeout

        future = Future()  # Never completes
        req = _make_req(json_schema="pending")
        req.grammar = future
        req.grammar_key = ("json", "pending")
        req.grammar_wait_ct = 0
        mgr.grammar_queue.append(req)

        result = mgr.get_ready_grammar_requests()
        self.assertEqual(len(result), 0)
        self.assertEqual(len(mgr.grammar_queue), 1)
        self.assertEqual(req.grammar_wait_ct, 1)

    def test_mixed_ready_and_pending(self):
        mgr = self._make_mgr()
        mgr.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = 100

        # Ready request
        grammar_obj = MagicMock(spec=BaseGrammarObject)
        grammar_obj.copy.return_value = grammar_obj
        done_future = Future()
        done_future.set_result(grammar_obj)
        ready_req = _make_req(json_schema="ready", rid="r1")
        ready_req.grammar = done_future
        ready_req.grammar_key = ("json", "ready")

        # Pending request
        pending_future = Future()
        pending_req = _make_req(json_schema="pending", rid="r2")
        pending_req.grammar = pending_future
        pending_req.grammar_key = ("json", "pending")
        pending_req.grammar_wait_ct = 0

        mgr.grammar_queue = [ready_req, pending_req]

        result = mgr.get_ready_grammar_requests()
        self.assertEqual(len(result), 1)
        self.assertIs(result[0], ready_req)
        self.assertEqual(len(mgr.grammar_queue), 1)
        self.assertIs(mgr.grammar_queue[0], pending_req)

    def test_empty_queue(self):
        """get_ready_grammar_requests on empty queue should return empty list."""
        mgr = self._make_mgr()
        result = mgr.get_ready_grammar_requests()
        self.assertEqual(len(result), 0)
        self.assertEqual(len(mgr.grammar_queue), 0)

    def test_progressive_timeout(self):
        """Request with partial wait_ct should timeout after remaining iterations."""
        mgr = self._make_mgr()
        mgr.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = 3

        future = Future()  # Never completes
        req = _make_req(json_schema="slow")
        req.grammar = future
        req.grammar_key = ("json", "slow")
        req.grammar_wait_ct = 2  # Already waited 2 iterations
        mgr.grammar_queue.append(req)

        # wait_ct increments to 3 (== max), should timeout
        result = mgr.get_ready_grammar_requests()
        self.assertEqual(len(result), 1)
        req.set_finish_with_abort.assert_called_once()
        self.assertIn("timed out", req.set_finish_with_abort.call_args[0][0])

    def test_ready_future_applies_request_budget_without_polluting_cache(self):
        mgr = self._make_mgr()

        grammar_obj = ReasonerGrammarObject(
            grammar=None, think_end_id=0, max_think_tokens=99
        )
        future = Future()
        future.set_result(grammar_obj)

        req = _make_req(json_schema="schema", custom_params={"thinking_budget": 4})
        req.grammar = future
        req.grammar_key = ("json", "schema")
        mgr.grammar_queue.append(req)

        result = mgr.get_ready_grammar_requests()

        self.assertEqual(len(result), 1)
        self.assertEqual(req.grammar.max_think_tokens, 4)
        cached_key, cached_value = mgr.grammar_backend.set_cache.call_args[0]
        self.assertEqual(cached_key, ("json", "schema"))
        self.assertEqual(cached_value.max_think_tokens, 99)

    @patch("sglang.srt.constrained.grammar_manager.torch.distributed.all_gather_object")
    def test_multi_rank_sync_intersects_ready_unions_failed(self, mock_all_gather):
        """With multiple ranks, ready = intersection, failed = union."""
        mgr = self._make_mgr()
        mgr.grammar_sync_size = 2  # Enable multi-rank path

        # Two requests: idx 0 ready on both ranks, idx 1 ready only on rank 0
        grammar_obj = MagicMock(spec=BaseGrammarObject)
        grammar_obj.copy.return_value = grammar_obj
        done_future = Future()
        done_future.set_result(grammar_obj)

        req0 = _make_req(json_schema="s0", rid="r0")
        req0.grammar = done_future
        req0.grammar_key = ("json", "s0")

        pending_future = Future()
        req1 = _make_req(json_schema="s1", rid="r1")
        req1.grammar = pending_future
        req1.grammar_key = ("json", "s1")
        req1.grammar_wait_ct = 0

        mgr.grammar_queue = [req0, req1]
        mgr.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = 100

        # Simulate all_gather: rank 0 has {0} ready, rank 1 has {0,1} ready
        def fake_all_gather(output_list, _obj, group=None):  # noqa: ARG001
            output_list[0] = ({0}, set())  # rank 0: only idx 0 ready
            output_list[1] = ({0, 1}, set())  # rank 1: both ready

        mock_all_gather.side_effect = fake_all_gather

        result = mgr.get_ready_grammar_requests()
        # Intersection of ready: {0} ∩ {0,1} = {0}
        self.assertEqual(len(result), 1)
        self.assertIs(result[0], req0)
        # req1 stays in queue
        self.assertEqual(len(mgr.grammar_queue), 1)
        self.assertIs(mgr.grammar_queue[0], req1)

    @patch("sglang.srt.constrained.grammar_manager.torch.distributed.all_gather_object")
    def test_multi_rank_sync_unions_failed(self, mock_all_gather):
        """Failed requests from any rank should be unioned."""
        mgr = self._make_mgr()
        mgr.grammar_sync_size = 2
        mgr.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = 1

        pending_future = Future()  # Never completes
        req = _make_req(json_schema="slow", rid="r0")
        req.grammar = pending_future
        req.grammar_key = ("json", "slow")
        req.grammar_wait_ct = 0

        mgr.grammar_queue = [req]

        # Simulate: rank 0 has no ready and idx 0 failed, rank 1 has no ready/failed
        def fake_all_gather(output_list, _obj, group=None):  # noqa: ARG001
            output_list[0] = (set(), {0})  # rank 0: idx 0 timed out
            output_list[1] = (set(), set())  # rank 1: nothing

        mock_all_gather.side_effect = fake_all_gather

        result = mgr.get_ready_grammar_requests()
        # Union of failed: {} ∪ {0} = {0}
        self.assertEqual(len(result), 1)
        req.set_finish_with_abort.assert_called_once()
        self.assertIn("timed out", req.set_finish_with_abort.call_args[0][0])
        self.assertEqual(len(mgr.grammar_queue), 0)


class _FakePPSendWork:
    def __init__(self):
        self.waited = False
        self.work = self

    def wait(self):
        self.waited = True


class _FakePPGroup:
    def __init__(self, recv_data=None):
        self.recv_data = recv_data
        self.recv_calls = []
        self.send_calls = []

    def recv_object(self, *, src, tag):
        self.recv_calls.append((src, tag))
        return self.recv_data

    def send_object(self, data, *, dst, async_send, tag):
        self.send_calls.append((data, dst, async_send, tag))
        return [_FakePPSendWork()]


class TestGrammarManagerPPSync(unittest.TestCase):
    """Test PP synchronization of grammar ready/failed indexes."""

    def _make_mgr_for_pp(self, pp_rank, pp_size, pp_group):
        scheduler = _make_scheduler()
        scheduler.server_args.skip_tokenizer_init = True
        scheduler.ps.pp_rank = pp_rank
        scheduler.ps.pp_size = pp_size
        scheduler.pp_group = pp_group
        mgr = GrammarManager(scheduler)
        mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend)
        return mgr

    def test_pp0_sends_ready_failed_without_recv(self):
        pp_group = _FakePPGroup()
        mgr = self._make_mgr_for_pp(pp_rank=0, pp_size=3, pp_group=pp_group)

        data = mgr._pp_sync_ready_failed({1}, {3})

        self.assertEqual(data, ({1}, {3}))
        self.assertEqual(pp_group.recv_calls, [])
        self.assertEqual(
            pp_group.send_calls,
            [(({1}, {3}), 1, True, P2PTag.GRAMMAR_PP_SYNC)],
        )

    def test_middle_pp_rank_receives_and_forwards_pp0_result(self):
        pp0_data = ({1, 2}, {4})
        pp_group = _FakePPGroup(recv_data=pp0_data)
        mgr = self._make_mgr_for_pp(pp_rank=1, pp_size=3, pp_group=pp_group)

        data = mgr._pp_sync_ready_failed(set(), set())

        self.assertEqual(data, pp0_data)
        self.assertEqual(pp_group.recv_calls, [(0, P2PTag.GRAMMAR_PP_SYNC)])
        self.assertEqual(
            pp_group.send_calls,
            [(pp0_data, 2, True, P2PTag.GRAMMAR_PP_SYNC)],
        )

    def test_last_pp_rank_receives_without_forwarding(self):
        pp0_data = ({0}, {2})
        pp_group = _FakePPGroup(recv_data=pp0_data)
        mgr = self._make_mgr_for_pp(pp_rank=2, pp_size=3, pp_group=pp_group)

        data = mgr._pp_sync_ready_failed(set(), set())

        self.assertEqual(data, pp0_data)
        self.assertEqual(pp_group.recv_calls, [(1, P2PTag.GRAMMAR_PP_SYNC)])
        self.assertEqual(pp_group.send_calls, [])

    def test_pp_sync_drains_previous_async_send_work(self):
        pp_group = _FakePPGroup()
        mgr = self._make_mgr_for_pp(pp_rank=0, pp_size=2, pp_group=pp_group)
        work = _FakePPSendWork()
        mgr.grammar_pp_sync_work_list = [work]

        mgr._pp_sync_ready_failed({1}, set())

        self.assertTrue(work.waited)


class TestStrictReasoningPaths(unittest.TestCase):
    """Test _enable_strict_thinking code paths in GrammarManager."""

    def _make_mgr(self):
        scheduler = _make_scheduler()
        scheduler.server_args.skip_tokenizer_init = True
        mgr = GrammarManager(scheduler)
        mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend)
        mgr._enable_strict_thinking = True
        return mgr

    def test_strict_unconstrained_request_gets_strict_grammar(self):
        """Request without json_schema/regex/ebnf should get strict-only grammar."""
        mgr = self._make_mgr()
        grammar_obj = MagicMock()
        mgr.grammar_backend.init_strict_reasoning_grammar.return_value = grammar_obj

        req = _make_req()  # No constraint
        req.require_reasoning = True
        result = mgr.process_req_with_grammar(req)

        self.assertFalse(result)  # Not added to grammar queue
        self.assertIs(req.grammar, grammar_obj)
        mgr.grammar_backend.init_strict_reasoning_grammar.assert_called_once_with(True)

    def test_strict_unconstrained_no_reasoning_flag(self):
        """Unconstrained request with require_reasoning=False still gets strict grammar."""
        mgr = self._make_mgr()
        grammar_obj = MagicMock()
        mgr.grammar_backend.init_strict_reasoning_grammar.return_value = grammar_obj

        req = _make_req()
        req.require_reasoning = False
        mgr.process_req_with_grammar(req)

        self.assertIs(req.grammar, grammar_obj)
        mgr.grammar_backend.init_strict_reasoning_grammar.assert_called_once_with(False)

    def test_strict_unconstrained_none_grammar_is_fine(self):
        """If init_strict_reasoning_grammar returns None, req.grammar stays None."""
        mgr = self._make_mgr()
        mgr.grammar_backend.init_strict_reasoning_grammar.return_value = None

        req = _make_req()
        req.require_reasoning = True
        mgr.process_req_with_grammar(req)

        self.assertIsNone(req.grammar)

    def test_strict_constrained_request_uses_normal_dispatch(self):
        """Request with json_schema should go through normal dispatch, not strict path."""
        mgr = self._make_mgr()
        future = MagicMock(spec=Future)
        mgr.grammar_backend.get_cached_or_future_value.return_value = (future, False)

        req = _make_req(json_schema='{"type": "object"}')
        req.require_reasoning = True
        result = mgr.process_req_with_grammar(req)

        self.assertTrue(result)  # Added to grammar queue
        mgr.grammar_backend.init_strict_reasoning_grammar.assert_not_called()

    def test_strict_not_set_skips_strict_path(self):
        """When _enable_strict_thinking=False, unconstrained requests get no grammar."""
        mgr = self._make_mgr()
        mgr._enable_strict_thinking = False

        req = _make_req()
        req.require_reasoning = True
        mgr.process_req_with_grammar(req)

        self.assertIsNone(req.grammar)
        mgr.grammar_backend.init_strict_reasoning_grammar.assert_not_called()

    def test_future_exception_creates_invalid_grammar(self):
        """Future.result() raising should create InvalidGrammarObject, not crash."""
        mgr = self._make_mgr()

        future = Future()
        future.set_exception(RuntimeError("compilation failed"))

        req = _make_req(json_schema='{"type": "object"}')
        req.require_reasoning = True
        req.grammar = future
        req.grammar_key = ("json", '{"type": "object"}')
        mgr.grammar_queue.append(req)

        mgr.SGLANG_GRAMMAR_POLL_INTERVAL = 0.001
        result = mgr.get_ready_grammar_requests()

        self.assertEqual(len(result), 1)
        self.assertIsInstance(result[0].grammar, InvalidGrammarObject)
        req.set_finish_with_abort.assert_called_once()


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