#!/usr/bin/env python3
"""Patch vLLM structured-output reasoning-boundary bugs.

Ports upstream vLLM PRs #44297 and #44993 to the b12x fork's vLLM source.

Background
----------
When a reasoning parser (e.g. glm45 -> Glm47MoeParser) emits the reasoning-end
marker (</think>) under speculative decoding (MTP, num_speculative_tokens=4),
the grammar FSM (used by response_format / guided decoding) fails to advance
across the reasoning->content boundary. This is vLLM bug #34650, fixed
upstream by:

  #44297 - Constrain bitmask and trim grammar advance at the reasoning boundary
  #44993 - Advance grammar across reasoning boundary (supersedes #44993's
           earlier conservative version; final merged form)

Without this patch, models emit prose or markdown-fence prefixes before JSON
in the content channel (e.g. ```{..." or "Here{..."), causing structured
output validation failures. The tool-eval-bench Structured Output category
scores 50% (6/12) instead of 100%.

The fix
-------
1. should_advance() accepts an optional new_token_ids argument and uses it as
   the authoritative step delta (instead of reconstructing the window from
   num_computed_tokens - num_output_placeholders, which skips the marker
   under async scheduling + spec-decode draft rejection).

2. At the boundary, should_advance() records reasoning_end_token_index and
   returns True for ALL constraint types (not just STRUCTURAL_TAG under spec
   decode). The scheduler then trims reasoning content before accept_tokens.

3. grammar_bitmask() detects a mid-window reasoning-end marker and constrains
   subsequent draft positions via the bitmask, skipping grammar advance
   through the marker itself.

4. The scheduler's update_from_output() passes new_token_ids to should_advance
   and calls trim_reasoning_for_advance() before accept_tokens.

Fixes: #34650, #43388, #48228.
"""
import sys
from pathlib import Path

VLLM_DIR = Path("/opt/venv/lib/python3.12/site-packages/vllm")


def patch_request():
    """Add reasoning_end_token_index field to StructuredOutputRequest."""
    path = VLLM_DIR / "v1" / "structured_output" / "request.py"
    content = path.read_text()
    if "reasoning_end_token_index" in content:
        print(f"  {path.name}: already patched, skipping")
        return
    old = "    reasoning_ended: bool | None = None\n"
    new = (
        "    reasoning_ended: bool | None = None\n"
        "    # Absolute index into the request's all_token_ids of the last\n"
        "    # reasoning token (the reasoning-end marker). Tokens at or before\n"
        "    # this index are reasoning content and must never be fed to the\n"
        "    # grammar. Only set when reasoning ends in a step whose tokens\n"
        "    # the scheduler advances immediately (structural tags +\n"
        "    # speculative decoding, see #42452).\n"
        "    reasoning_end_token_index: int | None = None\n"
    )
    if old not in content:
        print(f"  {path.name}: ERROR - could not find reasoning_ended field",
              file=sys.stderr)
        sys.exit(1)
    content = content.replace(old, new, 1)
    path.write_text(content)
    print(f"  {path.name}: patched")


def patch_init():
    """Patch __init__.py: imports, grammar_bitmask, should_advance, helpers."""
    path = VLLM_DIR / "v1" / "structured_output" / "__init__.py"
    content = path.read_text()
    if "trim_reasoning_for_advance" in content:
        print(f"  {path.name}: already patched, skipping")
        return

    # 1. Add Sequence to collections.abc import
    old_import = "from collections.abc import Iterable\n"
    new_import = "from collections.abc import Iterable, Sequence\n"
    if old_import not in content:
        print(f"  {path.name}: ERROR - could not find Iterable import",
              file=sys.stderr)
        sys.exit(1)
    content = content.replace(old_import, new_import, 1)

    # 2. Remove StructuredOutputOptions from backend_types import
    old_bo = (
        "from vllm.v1.structured_output.backend_types import (\n"
        "    StructuredOutputBackend,\n"
        "    StructuredOutputGrammar,\n"
        "    StructuredOutputOptions,\n"
        ")"
    )
    new_bo = (
        "from vllm.v1.structured_output.backend_types import (\n"
        "    StructuredOutputBackend,\n"
        "    StructuredOutputGrammar,\n"
        ")"
    )
    if old_bo not in content:
        print(f"  {path.name}: ERROR - could not find backend_types import",
              file=sys.stderr)
        sys.exit(1)
    content = content.replace(old_bo, new_bo, 1)

    # 3. Remove logger import and initialization (unused after should_advance
    #    no longer references StructuredOutputOptions)
    old_logger_import = "from vllm.logger import init_logger\n"
    if old_logger_import in content:
        content = content.replace(old_logger_import, "", 1)
    old_logger_init = "logger = init_logger(__name__)\n\n\n"
    if old_logger_init in content:
        content = content.replace(old_logger_init, "", 1)

    # 4. Replace grammar_bitmask serial-filling loop
    old_bitmask = '''                state_advancements = 0
                req_tokens = scheduled_spec_decode_tokens.get(req_id, ())
                if self.vllm_config.model_config.is_diffusion and req_tokens:
                    # Diffusion LLMs don't sample a bonus token after the
                    # scheduled positions, so don't append the -1 placeholder.
                    token_iter: Iterable[int] = req_tokens
                else:
                    token_iter = itertools.chain(req_tokens, (-1,))
                for token in token_iter:
                    self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),))
                    if token == -1:
                        # Stop advancing the grammar once we hit a padding token.
                        apply_bitmask = False
                    if apply_bitmask and not grammar.is_terminated():
                        accepted = grammar.accept_tokens(req_id, [token])
                        assert accepted, (token, req_id, scheduled_spec_decode_tokens)
                        state_advancements += 1
                    cumulative_index += 1
                if state_advancements > 0:
                    grammar.rollback(state_advancements)'''

    new_bitmask = '''                reasoner = self._get_reasoner(request)
                detect_reasoning_end = (
                    not apply_bitmask
                    and reasoner is not None
                    and not self.enable_in_reasoning
                )
                simulated_buf: list[int] | None = None
                history_len = 0

                state_advancements = 0
                post_reasoning_end_in_window = False
                req_tokens = scheduled_spec_decode_tokens.get(req_id, ())
                for i, token in enumerate(req_tokens):
                    self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),))
                    advance_grammar = apply_bitmask
                    if token == -1:
                        apply_bitmask = False
                        advance_grammar = False
                    elif (
                        detect_reasoning_end
                        and reasoner is not None
                        and not apply_bitmask
                    ):
                        if simulated_buf is None:
                            history = list(request.all_token_ids)
                            history_len = len(history)
                            simulated_buf = history + list(req_tokens)
                        simulated = simulated_buf[: history_len + i + 1]
                        if reasoner.is_reasoning_end_streaming(simulated, [token]):
                            # Reasoning ended mid-window. Constrain the rest
                            # of the window via bitmask. Skip grammar advance
                            # through the marker (it is reasoning content);
                            # try to advance through subsequent drafts so the
                            # next bitmask row reflects the post-advance state,
                            # but tolerate rejection since those drafts predate
                            # the bitmask and are not guaranteed valid.
                            apply_bitmask = True
                            advance_grammar = False
                            post_reasoning_end_in_window = True
                    if advance_grammar and not grammar.is_terminated():
                        accepted = grammar.accept_tokens(req_id, [token])
                        if accepted:
                            state_advancements += 1
                        elif not post_reasoning_end_in_window:
                            raise AssertionError(
                                (token, req_id, scheduled_spec_decode_tokens)
                            )
                    cumulative_index += 1
                # Diffusion LLMs don't sample a bonus token after the
                # scheduled positions, so skip its bitmask in that case.
                if not (self.vllm_config.model_config.is_diffusion and req_tokens):
                    # bonus_apply must be True when the bonus-row position
                    # should be grammar-constrained. Two triggers:
                    # - should_fill_bitmask(request): reasoning was already
                    #   over at step start (or no reasoner /
                    #   enable_in_reasoning).
                    # - apply_bitmask: reasoning ended mid-window in this
                    #   call and was flipped True after the marker;
                    #   should_fill_bitmask still returns False here because
                    #   reasoning_ended is only persisted later by
                    #   should_advance.
                    bonus_apply = self.should_fill_bitmask(request) or apply_bitmask
                    self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),))
                    cumulative_index += 1
                if state_advancements > 0:
                    grammar.rollback(state_advancements)'''

    if old_bitmask not in content:
        print(f"  {path.name}: ERROR - could not find grammar_bitmask serial loop",
              file=sys.stderr)
        sys.exit(1)
    content = content.replace(old_bitmask, new_bitmask, 1)

    # 5. Replace should_advance and add helper methods
    old_should_advance = '''    def should_advance(self, request: "Request") -> bool:
        if not request.use_structured_output:
            return False

        # To determine whether we can advance the FSM.
        # Supports thinking usage where we skip the reasoning components.
        if TYPE_CHECKING:
            assert request.structured_output_request is not None
            assert request.structured_output_request.grammar is not None
        # by default, we should always advance
        # for cases that don't use thinking mode.
        reasoner = self._get_reasoner(request)
        if reasoner is None:
            return True

        # if the model needs structured in reasoning, we should advance
        if self.enable_in_reasoning:
            return True

        structured_req = request.structured_output_request
        if structured_req.reasoning_ended:
            return True

        # Check if reasoning ends in *this* step
        delta_from = request.num_computed_tokens - request.num_output_placeholders
        all_token_ids = request.all_token_ids
        start = (
            delta_from if delta_from >= 0 else max(len(all_token_ids) + delta_from, 0)
        )
        if reasoner.is_reasoning_end_streaming(
            all_token_ids, itertools.islice(all_token_ids, start, None)
        ):
            structured_req.reasoning_ended = True

            # Reasoning just ended this step. Defer FSM advance until the next
            # pass (see reasoning_ended check above) for JSON/regex/choice/grammar:
            # advancing on the closing boundary token can accept tokens that still
            # belong to the reasoning stream. Structural tags are the only safe
            # same-step exception: they model phased output (e.g. thinking tag ->
            # answer tag), and speculative decoding must run grammar.validate_tokens
            # on draft tokens produced immediately after that transition.
            if (
                self.vllm_config.speculative_config is not None
                and structured_req.structured_output_key[0]
                == StructuredOutputOptions.STRUCTURAL_TAG
            ):
                return True

        return False'''

    new_should_advance = '''    def should_advance(
        self,
        request: "Request",
        new_token_ids: list[int] | None = None,
    ) -> bool:
        if not request.use_structured_output:
            return False

        # To determine whether we can advance the FSM.
        # Supports thinking usage where we skip the reasoning components.
        if TYPE_CHECKING:
            assert request.structured_output_request is not None
            assert request.structured_output_request.grammar is not None
        # by default, we should always advance
        # for cases that don't use thinking mode.
        reasoner = self._get_reasoner(request)
        if reasoner is None:
            return True

        # if the model needs structured in reasoning, we should advance
        if self.enable_in_reasoning:
            return True

        structured_req = request.structured_output_request
        if structured_req.reasoning_ended:
            return True

        # Check if reasoning ends in *this* step.
        # When the caller passes new_token_ids (the tokens that were just
        # appended this step), use it directly as the delta window. The
        # placeholder-derived fallback assumes num_output_placeholders ==
        # len(new_token_ids), which breaks under async scheduling + spec
        # decode when some drafts are rejected (#43388): the placeholder
        # count remains > 0 after the step and the computed delta window
        # starts past the reasoning-end marker.
        all_token_ids = request.all_token_ids
        if new_token_ids:
            # The tokens were already appended this step, so the step window
            # starts exactly len(new_token_ids) from the end.
            start = len(all_token_ids) - len(new_token_ids)
            delta_ids: Iterable[int] = new_token_ids
        else:
            delta_from = request.num_computed_tokens - request.num_output_placeholders
            start = (
                delta_from
                if delta_from >= 0
                else max(len(all_token_ids) + delta_from, 0)
            )
            delta_ids = itertools.islice(all_token_ids, start, None)
        if reasoner.is_reasoning_end_streaming(all_token_ids, delta_ids):
            structured_req.reasoning_ended = True

            # Record the boundary so the scheduler can exclude reasoning tokens.
            end_index = self._find_reasoning_end_index(reasoner, all_token_ids, start)

            structured_req.reasoning_end_token_index = end_index
            return True

        return False

    @staticmethod
    def _find_reasoning_end_index(
        reasoner: "ReasoningParser", all_token_ids: Sequence[int], start: int
    ) -> int:
        """Locates the last reasoning token within all_token_ids[start:].

        Returns the absolute index of the token at which
        is_reasoning_end_streaming first fires. Falls back to the final
        index when no single token triggers the detection (e.g. a
        multi-token marker only recognized on the full delta), which
        conservatively treats the whole step as reasoning content.
        """
        prefix = list(itertools.islice(all_token_ids, start))
        for idx in range(start, len(all_token_ids)):
            token = all_token_ids[idx]
            prefix.append(token)
            if reasoner.is_reasoning_end_streaming(prefix, [token]):
                return idx
        return len(all_token_ids) - 1

    def trim_reasoning_for_advance(
        self, request: "Request", new_token_ids: list[int]
    ) -> list[int]:
        """Drops reasoning content from tokens about to advance the grammar.

        When reasoning ends mid-step (see should_advance), the step's output
        still contains reasoning tokens up to and including the end marker.
        Those are not grammar content: feeding them to accept_tokens makes
        the grammar reject the marker and kills the request (#44006).

        Returns the suffix of new_token_ids that follows the reasoning-end
        marker. Steps fully after the boundary are returned unchanged.
        """
        structured_req = request.structured_output_request
        if structured_req is None:
            return new_token_ids
        end_idx = structured_req.reasoning_end_token_index
        if end_idx is None:
            return new_token_ids
        first_idx = len(request.all_token_ids) - len(new_token_ids)
        num_reasoning = end_idx + 1 - first_idx
        if num_reasoning <= 0:
            return new_token_ids
        return new_token_ids[num_reasoning:]'''

    if old_should_advance not in content:
        print(f"  {path.name}: ERROR - could not find should_advance method",
              file=sys.stderr)
        sys.exit(1)
    content = content.replace(old_should_advance, new_should_advance, 1)

    path.write_text(content)
    print(f"  {path.name}: patched")


def patch_scheduler():
    """Patch update_from_output to pass new_token_ids and trim before accept."""
    path = VLLM_DIR / "v1" / "core" / "sched" / "scheduler.py"
    content = path.read_text()
    if "trim_reasoning_for_advance" in content:
        print(f"  {path.name}: already patched, skipping")
        return

    old = '''            if new_token_ids and self.structured_output_manager.should_advance(request):
                struct_output_request = request.structured_output_request
                assert struct_output_request is not None
                assert struct_output_request.grammar is not None
                if not struct_output_request.grammar.accept_tokens(  # type: ignore[union-attr]
                    req_id, new_token_ids
                ):
                    logger.error(
                        "Unexpected: grammar rejected tokens %s for request %s. "
                        "Terminating request.",
                        new_token_ids,
                        req_id,
                    )
                    request.status = RequestStatus.FINISHED_ERROR
                    request.resumable = False
                    stopped = True'''

    new = '''            if new_token_ids and self.structured_output_manager.should_advance(
                request, new_token_ids=new_token_ids
            ):
                struct_output_request = request.structured_output_request
                assert struct_output_request is not None
                grammar = struct_output_request.grammar
                assert grammar is not None
                # new_token_ids can be a mixed block of reasoning content, then
                # the reasoning end marker, then the start of the grammar content.
                # Trim the reasoning content so the grammar only sees grammar content.
                advance_token_ids = (
                    self.structured_output_manager.trim_reasoning_for_advance(
                        request, new_token_ids
                    )
                )
                if advance_token_ids and not grammar.accept_tokens(
                    req_id, advance_token_ids
                ):
                    logger.error(
                        "Unexpected: grammar rejected tokens %s for request %s. "
                        "Terminating request.",
                        advance_token_ids,
                        req_id,
                    )
                    request.status = RequestStatus.FINISHED_ERROR
                    request.resumable = False
                    stopped = True'''

    if old not in content:
        print(f"  {path.name}: ERROR - could not find should_advance block",
              file=sys.stderr)
        sys.exit(1)
    content = content.replace(old, new, 1)
    path.write_text(content)
    print(f"  {path.name}: patched")


if __name__ == "__main__":
    print("Applying structured output reasoning-boundary patch...")
    print(f"  vLLM dir: {VLLM_DIR}")
    patch_request()
    patch_init()
    patch_scheduler()
    print("Done.")