"""mlx_lm.server, plus structured output, reply limits, and cancellation.

Started by `agency local serve <model>`. It is `mlx_lm.server` with three
additions.

A request whose `response_format` names a JSON schema gets a reply that fits
the schema, because the schema is enforced while the reply is generated.
`mlx_lm.server` on its own never reads `response_format`, so a typed `llm()`
call on an MLX model came back as prose.

A reply that has stopped making progress is cut short. Three limits watch
every reply as it is written. A thinking model may spend at most a budget of
tokens thinking, after which it has to close its thinking and answer. A reply
that keeps second-guessing itself ("But wait", "Hmm", "Let me reconsider")
past a count, or repeats a whole sentence a few times, is cut short too. The
last two watch the thinking by default and the answer only when asked,
because an answer repeats itself for honest reasons (a refrain, a table, the
same line of code three times) and thinking rarely does. Without these, a
small model that talks itself in circles runs to `max_tokens`, which on a
large model is minutes of the machine for nothing.

A reply whose client has gone is stopped, where `mlx_lm.server` would finish
it for nobody.

How the reply is constrained. A model writes one token at a time, scoring
every token it knows at each step. The grammar library llguidance knows the
schema and how much of the JSON has been written, and hands back a bitmask of
the tokens that could come next. `Watcher` below is a logits processor: at
each step it sets the score of every token outside that mask to minus
infinity, so the sampler cannot pick it. Once the JSON is complete, only the
end-of-reply token is left.

A thinking model is left alone while it thinks. A reply goes through up to
three phases:

  start      the model may open a `<think>` block, or start the answer
  reasoning  anything goes, except ending the reply, until `</think>`
  answer     every token must fit the schema, when there is one

When the chat template already opened a `<think>` block at the end of the
prompt, the reply starts in the reasoning phase. mlx_lm works that out for
its own purposes in `_tokenize`, and this script reads its answer. Once the
model closes the block, the answer starts.

How a limit cuts a reply short depends on the phase. While thinking, the
only tokens allowed are the ones that close the thinking block, so the model
has to answer. In an answer with a schema, the only tokens allowed for one
step are the ones that close the JSON string being written, so the field
ends and the schema carries the reply to its end. In a plain answer, only
the end-of-reply token is allowed. An answer cut short is reported with
`finish_reason: "length"`, the way a reply that hit `max_tokens` is, so the
caller can tell it from a complete one.

Each limit can be set per request, in the body: `reasoning_budget`,
`hedge_limit`, and `repeat_limit`, each a count with 0 meaning no limit, and
`limit_answers`, true to watch the answer as well as the thinking. A request
that names none gets the server's flags, `--reasoning-budget`,
`--hedge-limit`, `--repeat-limit`, and `--limit-answers`. The reasoning
budget defaults to half the request's `max_tokens`, and is always held far
enough under `max_tokens` to leave room for the answer.

How cancellation works. A client that gives up on a reply (Agency's per-call
timeout, or Ctrl-C in the run) closes its socket. `mlx_lm.server` only
touches that socket once the whole reply is generated, so it never notices;
the reply runs on to `max_tokens`, holding its share of the GPU and its
cache, and every other request in the batch runs slower for it. The handler
thread that waits for tokens here looks at its socket every half second
instead, and when the socket has been closed it tells the generator to drop
the reply. mlx_lm drops it at the next point it looks: on the batched path
that is after the next chunk of prompt or the next token, so within a
moment; on the single path, which a draft model or a seed puts a request
on, only once the whole prompt has been read, so a long prompt is still read
to the end. The batch frees the reply's cache on its next step.

With a draft model (`--draft-model`), mlx_lm generates speculatively: the
draft guesses several tokens, the main model checks them in one pass, and
the guesses it rejects are taken back. The logits processor is called for
the guesses too, and then sees the history shrink. `Watcher` keeps what it
learned from each token, so it can take back exactly the tokens that were
rejected, including what it told the grammar.

A request that carries tools keeps its `response_format` but the schema is
not enforced, because a tool call is not JSON and the constraint would break
it. That matches what the llama.cpp backend does. The limits still apply.
The exception is a request that also sends `tool_choice: "none"`: the caller
has ruled a tool call out, so the schema is enforced.

This script reaches into `mlx_lm.server`: it subclasses its request handler
and response generator, and replaces two module-level names, `run` and
`_make_logits_processors`. `agency local serve` pins the mlx-lm version for
that reason; a new version needs these seams checked.

Options are `mlx_lm.server`'s own (--model, --host, --port, --max-tokens, and
the rest) plus the limit flags above. Distributed serving (`mlx.distributed`)
is not supported; the script always runs the HTTP server in its own process.
"""

import argparse
import json
import logging
import os
import re
import select
import socket
import sys
import time
from collections import Counter
from dataclasses import dataclass, replace
from queue import Empty as QueueEmpty
from queue import Queue

import numpy as np
from llguidance import LLMatcher
from llguidance.hf import from_tokenizer
from llguidance.mlx import apply_token_bitmask
from llguidance.numpy import allocate_token_bitmask, fill_next_token_bitmask
import mlx.core as mx
from mlx_lm import server
from mlx_lm.generate import GenerationBatch
from mlx_lm.models.cache import can_trim_prompt_cache


def schema_can_be_enforced(tools, tool_choice):
    """Whether a reply can be held to its `response_format`. It cannot while
    the model may call a tool, since a tool call is not JSON. A request with
    no tools, or one whose `tool_choice` is "none", rules a tool call out."""
    return not tools or tool_choice == "none"


def grammar_for(response_format):
    """The llguidance grammar for a request's `response_format`, or None when
    it asks for none. Raises ValueError for a shape this script cannot
    honour, which the handler turns into a 400."""
    if response_format is None:
        return None
    if not isinstance(response_format, dict):
        raise ValueError("response_format must be an object.")
    kind = response_format.get("type")
    if kind in (None, "text"):
        return None
    if kind == "json_object":
        return LLMatcher.grammar_from_json_schema({"type": "object"})
    if kind == "json_schema":
        schema = response_format.get("json_schema")
        if not isinstance(schema, dict) or "schema" not in schema:
            raise ValueError("response_format.json_schema.schema is missing.")
        grammar = LLMatcher.grammar_from_json_schema(schema["schema"])
        problem = LLMatcher.validate_grammar(grammar)
        if problem and not LLMatcher.is_validate_warning(problem):
            raise ValueError(f"The schema cannot be enforced: {problem}")
        return grammar
    raise ValueError(f"response_format.type {kind!r} is not supported.")


# ---------------------------------------------------------------------------
# Token bitmasks: one bit per token, set when the token may come next.
# ---------------------------------------------------------------------------


def bit_index(token_id):
    return token_id // 32, np.uint32(1) << np.uint32(token_id % 32)


def allow(mask, token_id):
    word, bit = bit_index(token_id)
    mask.view(np.uint32)[0, word] |= bit


def forbid(mask, token_id):
    word, bit = bit_index(token_id)
    mask.view(np.uint32)[0, word] &= ~bit


def only(mask, token_ids):
    """Leave nothing allowed but `token_ids`."""
    mask.fill(0)
    for token_id in token_ids:
        allow(mask, token_id)


def intersect(mask, other):
    """Keep only the tokens both masks allow, in `mask`."""
    np.bitwise_and(mask, other, out=mask)


def allows_nothing(mask):
    return not mask.any()


# ---------------------------------------------------------------------------
# The limits, and how a reply is judged against them.
# ---------------------------------------------------------------------------


@dataclass
class Limits:
    """How far a reply may go before it is cut short. Each count's 0 means
    no limit. A reasoning budget of None means half the request's
    `max_tokens`, decided per request. `limit_answers` watches the answer
    for hedging and repeats too; by default only the thinking is watched."""

    reasoning_budget: int | None
    hedge_limit: int
    repeat_limit: int
    limit_answers: bool


# The server's flags. Replaced in __main__ from the command line.
SERVER_LIMITS = Limits(reasoning_budget=None, hedge_limit=12, repeat_limit=3, limit_answers=False)

# A reply is judged every this many tokens, not every token: the scans below
# read the text so far.
CHECK_EVERY = 16

# How much of the reply the hedge and repeat scans look at: the newest this
# many tokens. A loop says the same thing every few lines; an honest long
# reply says "wait" a dozen times over thousands of tokens. Counting over a
# window tells the two apart where a count over the whole reply cannot.
WINDOW = 2048

# The phrases a model uses when it doubts what it just wrote. "Wait" only
# counts with punctuation after it, so "wait until the lock is released"
# does not.
HEDGES = re.compile(
    r"\b(?:But wait|Wait)(?=[,.!?:;—–-])"
    r"|\bHmm+\b|\bHold on\b|\bActually,? wait\b"
    r"|\bLet me (?:reconsider|re-?check|double-?check|re-?examine|re-?read|verify|think again)\b"
)

SENTENCE_END = re.compile(r"(?<=[.!?])\s+|\n+")

# A sentence shorter than this is not counted as a repeat. "Yes." and "So
# the answer is 51." recur in honest working.
MIN_SENTENCE_WORDS = 6

# A token that ends the JSON string being written: a double quote first,
# after any spaces (which are part of the string).
STRING_CLOSER = re.compile(r'^\s*"')


@dataclass
class TokenMasks:
    """Two facts about every token of one model, as bitmasks: whether it
    can end a JSON string, and whether it is anything but whitespace."""

    closers: np.ndarray
    non_space: np.ndarray


def token_masks(hf_tokenizer, vocab_width):
    """Decodes the whole vocabulary once, which takes a moment."""
    closers = allocate_token_bitmask(1, vocab_width)
    closers.fill(0)
    non_space = allocate_token_bitmask(1, vocab_width)
    non_space.fill(0)
    texts = hf_tokenizer.batch_decode([[i] for i in range(len(hf_tokenizer))])
    for token_id, text in enumerate(texts):
        if STRING_CLOSER.match(text):
            allow(closers, token_id)
        if text.strip() != "":
            allow(non_space, token_id)
    return TokenMasks(closers=closers, non_space=non_space)


def count_hedges(text):
    return len(HEDGES.findall(text))


def most_repeated_sentence(text):
    """How many times the most repeated long sentence appears in `text`."""
    sentences = [" ".join(s.split()) for s in SENTENCE_END.split(text)]
    long_ones = [s for s in sentences if len(s.split()) >= MIN_SENTENCE_WORDS]
    if not long_ones:
        return 0
    return max(Counter(long_ones).values())


def answer_reserve(max_tokens):
    """Tokens kept free for the answer after the thinking: an eighth of the
    reply, and never fewer than 256."""
    return max(256, max_tokens // 8)


def limits_for(body, max_tokens, defaults):
    """The limits for one request: the fields it names, else the server's
    flags. The reasoning budget is held under `max_tokens` by the answer's
    reserve, whoever set it. Raises ValueError for a field of the wrong
    shape, which the handler turns into a 400."""

    def count(name, default):
        value = body.get(name)
        if value is None:
            return default
        if isinstance(value, bool) or not isinstance(value, int) or value < 0:
            raise ValueError(f"{name} must be a whole number, 0 or more.")
        return value

    limit_answers = body.get("limit_answers", defaults.limit_answers)
    if not isinstance(limit_answers, bool):
        raise ValueError("limit_answers must be true or false.")

    budget = count("reasoning_budget", defaults.reasoning_budget)
    if budget is None:
        budget = max_tokens // 2
    most = max(0, max_tokens - answer_reserve(max_tokens))
    if budget > most:
        logging.info(
            f"A reasoning budget of {budget} leaves no room to answer in {max_tokens} tokens; using {most}."
        )
        budget = most
    return Limits(
        reasoning_budget=budget,
        hedge_limit=count("hedge_limit", defaults.hedge_limit),
        repeat_limit=count("repeat_limit", defaults.repeat_limit),
        limit_answers=limit_answers,
    )


@dataclass
class WatcherState:
    """Everything about a reply that one more token can change. A snapshot
    is kept per token, so that a token taken back by speculative decoding
    takes its effects back with it."""

    phase: str
    reasoning_tokens: int
    # While the thinking is being closed, the index of the next token of
    # the closing sequence. None otherwise.
    closing: int | None
    # The token the judged text starts at. The answer is judged on its own,
    # not with the thinking before it, and a JSON field cut short does not
    # count against the fields after it.
    judged_from: int
    # A plain answer being cut short: only end-of-reply tokens may come.
    ending: bool
    # A JSON answer being cut short: for one step, only tokens that close
    # the string being written may come.
    closing_string: bool
    # After that, no whitespace-only tokens: JSON allows any amount of
    # whitespace between its parts, and a model whose string was closed
    # under it will pad with tabs to max_tokens if it is let.
    compact: bool
    # While the answer's header is being written (Harmony's final-channel
    # header, before the JSON), the index of its next token. None when
    # there is no header, or it has been written.
    header: int | None
    # In the answer, how many tokens of the thinking marker the reply has
    # just written: a format whose reply is several messages can open
    # thinking again after the answer began.
    opening: int
    # Whether the thinking was cut short for its budget. A reply that was
    # may not open thinking again.
    thought_cut: bool
    # How many tokens the grammar has been fed, for taking them back.
    consumed: int
    # The grammar rejected a token; the rest of the reply runs unconstrained.
    broken: bool
    # Why the answer was cut short, or None. Reported as finish_reason
    # "length". Closing the thinking is not a cut: the answer still comes.
    cut: str | None


class Watcher:
    """The logits processor for one reply. mlx_lm calls it with the token
    history so far and the scores for the next token, and takes back the
    scores it should sample from. It enforces the schema, when there is one,
    and the limits.

    The history it sees starts partway through the prompt, since mlx_lm
    processes most of the prompt without calling any processor, and grows
    by one token per step. The first call fixes where the reply begins;
    every later call feeds the tokens since then into the phase machine.
    Under speculative decoding the history can also shrink, when the main
    model rejects the draft's guesses; see `rewind`."""

    def __init__(
        self,
        tokenizer,
        ll_tokenizer_for,
        token_masks_for,
        piece_for,
        grammar,
        initial_state,
        limits,
        vocab_width,
    ):
        self.tokenizer = tokenizer
        self.ll_tokenizer_for = ll_tokenizer_for
        self.token_masks_for = token_masks_for
        self.piece_for = piece_for
        self.grammar = grammar
        self.limits = limits
        # The main model's output width. A draft model's logits may be
        # narrower, and the masks have to cover the main model's tokens.
        self.vocab_width = vocab_width
        if not tokenizer.has_thinking:
            phase = "answer"
        elif initial_state == "reasoning":
            phase = "reasoning"
        else:
            phase = "start"
        self.think_start = tokenizer.think_start_tokens
        self.think_end = tokenizer.think_end_tokens
        # The tokens a reply may write between its thinking and its answer,
        # when its format has some: Harmony's final-channel header, with the
        # turn marker in front of it after thinking and without at the
        # reply's start. They are let through in front of the JSON and not
        # fed to the grammar.
        self.answer_headers = [tuple(h) for h in getattr(tokenizer, "_harmony_headers", None) or []]
        self.state = WatcherState(
            phase=phase,
            reasoning_tokens=0,
            closing=None,
            judged_from=0,
            ending=False,
            closing_string=False,
            compact=False,
            consumed=0,
            broken=False,
            cut=None,
            header=None,
            opening=0,
            thought_cut=False,
        )
        self.initial = replace(self.state)
        self.seen = None
        # Where the reply starts in the history mlx_lm hands over.
        self.start = None
        # Per token of the reply: its id, its text, and the state after it.
        self.reply = []
        self.pieces = []
        self.states = []
        self.matcher = None
        self.mask = None
        self.free_mask = None
        self.end_mask = None
        self.token_masks = None

    @property
    def cut(self):
        return self.state.cut

    def prepare(self, vocab_width):
        """Built on the first call, because only the logits say how wide the
        model's output is. A padded output layer is wider than the tokenizer's
        vocabulary, and the mask has to cover the whole row."""
        self.mask = allocate_token_bitmask(1, vocab_width)
        # While thinking, everything is allowed but ending the reply. The
        # model has to close the block and write the answer.
        self.free_mask = allocate_token_bitmask(1, vocab_width)
        for token_id in self.tokenizer.eos_token_ids:
            forbid(self.free_mask, token_id)
        self.end_mask = allocate_token_bitmask(1, vocab_width)
        only(self.end_mask, self.tokenizer.eos_token_ids)
        if self.grammar is not None:
            ll_tokenizer = self.ll_tokenizer_for(vocab_width)
            self.matcher = LLMatcher(ll_tokenizer, self.grammar)
            if self.matcher.is_error():
                raise ValueError(self.matcher.get_error())
            self.token_masks = self.token_masks_for(vocab_width)

    def consume(self, token_id):
        """Feed one JSON token to the matcher. A token it rejects, which can
        only happen when llguidance and mlx_lm disagree about the tokenizer,
        leaves it in an error state for good; from then on the reply runs
        unconstrained, and the log says why, so the caller gets a reply that
        fails validation rather than one built from a broken matcher."""
        if self.matcher is None or self.state.broken:
            return
        if self.matcher.consume_token(token_id):
            self.state.consumed += 1
            return
        self.state.broken = True
        logging.error(
            f"The schema constraint failed and the rest of this reply is "
            f"unconstrained: {self.matcher.get_error()}"
        )

    def openers(self):
        """The token sequences a reply may start with before its answer:
        the thinking marker, and the answer headers when the format has
        them."""
        return [self.think_start] + self.answer_headers

    def headers_from(self, written):
        """The answer headers that begin with `written` and go on past it."""
        return [h for h in self.answer_headers if h[: len(written)] == written and len(h) > len(written)]

    def enter_answer(self, header_done=False):
        self.state.phase = "answer"
        self.state.closing = None
        self.state.judged_from = len(self.pieces)
        self.state.header = 0 if self.answer_headers and not header_done else None
        self.state.opening = 0

    def enter_reasoning(self):
        self.state.phase = "reasoning"
        self.state.closing = None
        self.state.opening = 0

    def watch_for_thinking(self):
        """In the answer, whether the reply has opened thinking again: a
        format whose reply is several messages can. The reply's tail is
        read, not the header or grammar state, because the tokens of a
        marker can start out looking like something else. Harmony's
        `<|start|>assistant<|channel|>analysis<|message|>` starts out
        looking like the final header, for one. A whole marker goes back
        to the reasoning phase, the tokens already thought still counted;
        part of one is kept in `opening`. A model without thinking has no
        marker to watch for."""
        if not self.think_start:
            return
        state = self.state
        width = len(self.think_start)
        if tuple(self.reply[-width:]) == self.think_start:
            self.enter_reasoning()
            return
        state.opening = 0
        for written in range(min(width - 1, len(self.reply)), 0, -1):
            if tuple(self.reply[-written:]) == self.think_start[:written]:
                state.opening = written
                return

    def step(self, token_id):
        """One generated token moves the phase machine. The answer is judged
        from its first token, so a reply that starts with the answer enters
        the phase before the token is added to the text, and one that closes
        its thinking enters it after, so the closing tokens are left out."""
        state = self.state
        self.reply.append(token_id)
        if state.phase == "start":
            # A marker can be several tokens (Harmony's are three), so the
            # reply stays at its start while it could still be one, and
            # the tokens so far are the answer's first ones when it turns
            # out to be neither.
            self.pieces.append(self.piece_for(token_id))
            so_far = tuple(self.reply)
            if so_far == self.think_start:
                self.enter_reasoning()
            elif so_far in self.answer_headers:
                self.enter_answer(header_done=True)
            elif not any(opener[: len(so_far)] == so_far for opener in self.openers()):
                self.enter_answer(header_done=True)
                state.judged_from = 0
                for earlier in so_far:
                    self.consume(earlier)
        elif state.phase == "reasoning":
            self.pieces.append(self.piece_for(token_id))
            state.reasoning_tokens += 1
            if state.closing is not None and token_id == self.think_end[state.closing]:
                state.closing += 1
            tail = tuple(self.reply[-len(self.think_end) :])
            if tail == self.think_end:
                self.enter_answer()
        elif state.header is not None and any(
            token_id == h[state.header]
            # The header's tokens before this one: the reply's tail, less
            # this token, which step() has already appended.
            for h in self.headers_from(tuple(self.reply[len(self.reply) - 1 - state.header : -1]))
        ):
            # A header is passed over, not judged and not fed to the grammar.
            self.pieces.append(self.piece_for(token_id))
            state.header += 1
            if tuple(self.reply[-state.header :]) in self.answer_headers:
                state.header = None
                state.judged_from = len(self.pieces)
            self.watch_for_thinking()
        else:
            state.header = None
            self.pieces.append(self.piece_for(token_id))
            self.consume(token_id)
            self.watch_for_thinking()
        if len(self.reply) % CHECK_EVERY == 0:
            self.check()
        self.states.append(replace(state))

    def rewind(self, count):
        """Take back the last `count` tokens: the draft's guesses the main
        model rejected. Everything they changed is restored from the
        snapshot before them, and the grammar takes back what it was fed."""
        fed = self.state.consumed
        del self.reply[-count:]
        del self.pieces[-count:]
        del self.states[-count:]
        self.state = replace(self.states[-1]) if self.states else replace(self.initial)
        if self.matcher is not None and fed > self.state.consumed:
            self.matcher.rollback(fed - self.state.consumed)

    def resync(self, tokens):
        """mlx_lm 0.31.3 only ever shortens the history to a prefix of itself
        and appends to it, so taking back tokens by count keeps the reply in
        step with it. This guards the day that changes: when the last token
        of the history is not the last token stepped, find where the two
        part, take back from there, and step the rest. One element per call
        when they agree."""
        if not self.reply or int(tokens[-1].item()) == self.reply[-1]:
            return
        history = tokens[self.start :].tolist()
        shared = min(len(history), len(self.reply))
        parted = next((i for i in range(shared) if history[i] != self.reply[i]), shared)
        logging.warning(
            f"The token history parted from the reply at token {parted}; reading it again from there."
        )
        self.rewind(len(self.reply) - parted)
        for token_id in history[parted:]:
            self.step(token_id)

    def judged_text(self):
        """The newest WINDOW tokens of the text being judged."""
        start = max(self.state.judged_from, len(self.pieces) - WINDOW)
        return "".join(self.pieces[start:])

    def check(self):
        """Has the reply gone past a limit? Called every CHECK_EVERY tokens."""
        state = self.state
        if state.phase == "reasoning":
            if state.closing is not None:
                return
            budget = self.limits.reasoning_budget
            if budget and state.reasoning_tokens >= budget:
                self.cut_short(f"has thought for {state.reasoning_tokens} tokens, its budget")
                return
        elif state.ending or not self.limits.limit_answers:
            return
        text = self.judged_text()
        if self.limits.hedge_limit:
            hedges = count_hedges(text)
            if hedges >= self.limits.hedge_limit:
                self.cut_short(f"has second-guessed itself {hedges} times")
                return
        if self.limits.repeat_limit:
            repeats = most_repeated_sentence(text)
            if repeats >= self.limits.repeat_limit:
                self.cut_short(f"has repeated a sentence {repeats} times")

    def cut_short(self, why):
        state = self.state
        if state.phase == "reasoning":
            logging.info(f"The reply {why}; closing its thinking.")
            state.closing = 0
            state.thought_cut = True
        elif self.matcher is not None:
            logging.info(f"The reply {why}; closing the JSON string it is writing.")
            state.closing_string = True
            state.compact = True
            state.judged_from = len(self.pieces)
            state.cut = why
        else:
            logging.info(f"The reply {why}; ending it.")
            state.ending = True
            state.cut = why

    def __call__(self, tokens, logits):
        count = tokens.shape[-1]
        if self.seen is None:
            self.prepare(max(logits.shape[-1], self.vocab_width))
            self.seen = count
            self.start = count
        else:
            if count < self.seen:
                self.rewind(self.seen - count)
            for token_id in tokens[self.seen :].tolist():
                self.step(token_id)
            self.seen = count
            self.resync(tokens)
        state = self.state
        if state.phase == "reasoning":
            if state.closing is not None:
                only(self.mask, [self.think_end[state.closing]])
                return apply_token_bitmask(logits, self.mask)
            return apply_token_bitmask(logits, self.free_mask)
        if state.ending:
            return apply_token_bitmask(logits, self.end_mask)
        if self.matcher is None or state.broken:
            # No schema, so the answer is free, with one exception: thinking
            # that was cut for its budget may not be opened again. The first
            # token that commits the reply to thinking is refused. For a
            # marker of one token, such as `<think>`, that is its only token,
            # refused everywhere in the answer. For a longer one it is the
            # second: Harmony's `<|channel|>` also starts the final and
            # commentary headers, so only `analysis` after it is refused, and
            # the model is left free to take one of those instead. Refusing a
            # later token would put part of the marker in the reply's text
            # and leave the model stuck partway through it.
            if (
                state.phase == "answer"
                and state.thought_cut
                and self.think_start
                and state.opening == min(1, len(self.think_start) - 1)
            ):
                self.mask.fill(-1)
                forbid(self.mask, self.think_start[state.opening])
                return apply_token_bitmask(logits, self.mask)
            return logits
        if state.phase == "start" and self.reply:
            # Part of a marker so far: only what completes one may come.
            so_far = tuple(self.reply)
            only(self.mask, [o[len(so_far)] for o in self.openers() if o[: len(so_far)] == so_far])
            return apply_token_bitmask(logits, self.mask)
        if state.phase == "answer" and state.header:
            written = tuple(self.reply[-state.header :])
            only(self.mask, [h[state.header] for h in self.headers_from(written)])
            return apply_token_bitmask(logits, self.mask)
        fill_next_token_bitmask(self.matcher, self.mask)
        if state.phase == "start":
            for opener in self.openers():
                allow(self.mask, opener[0])
        elif state.phase == "answer" and state.header == 0:
            for header in self.answer_headers:
                allow(self.mask, header[0])
        if state.closing_string:
            # The tokens that both fit the schema and close the string. None
            # of them is allowed mid-escape; then the next step tries again.
            self.narrow(self.token_masks.closers, once=True)
        elif state.compact:
            self.narrow(self.token_masks.non_space)
        return apply_token_bitmask(logits, self.mask)

    def narrow(self, to, once=False):
        """Keep only the tokens of `self.mask` that `to` allows as well, when
        that leaves any. With `once`, a narrowing that took ends the closing
        of the string."""
        narrowed = self.mask.copy()
        intersect(narrowed, to)
        if allows_nothing(narrowed):
            return
        np.copyto(self.mask, narrowed)
        if once:
            self.state.closing_string = False


# ---------------------------------------------------------------------------
# Cancellation.
# ---------------------------------------------------------------------------

# How often a handler looks at its socket while it waits for tokens.
CLIENT_POLL_SECONDS = 0.5

# What the queue gave back when it gave back nothing in time.
WAITING = object()


def client_gone(connection):
    """Whether the client closed its end of the socket. A client sends
    nothing after its request, so the socket is only readable once the
    client has closed it, and a peek then returns no bytes. A socket that
    errors is gone too."""
    try:
        readable, _, _ = select.select([connection], [], [], 0)
    except ValueError:
        # A file descriptor select cannot watch. Not knowing is not gone.
        return False
    except OSError:
        return True
    if not readable:
        return False
    try:
        return connection.recv(1, socket.MSG_PEEK) == b""
    except OSError:
        return True


# ---------------------------------------------------------------------------
# Harmony: the channel format gpt-oss writes.
# ---------------------------------------------------------------------------
#
# A gpt-oss reply is a series of messages, each on a channel and each
# ended by `<|end|>`, with `<|start|>assistant` opening the next. Thinking
# goes on the analysis channel, the answer on the final channel, and a
# tool call on the commentary channel, ended by `<|call|>` instead:
#
#   <|channel|>analysis<|message|>…thinking…<|end|><|start|>assistant
#   <|channel|>final<|message|>…answer…<|return|>
#
#   <|channel|>commentary to=functions.NAME <|constrain|>json<|message|>{…}<|call|>
#
# The recipient may go in the role header instead, which is how the
# model's own template renders past calls, so the model often writes it
# that way after a tool round:
#
#   <|start|>assistant to=functions.NAME<|channel|>commentary json<|message|>{…}<|call|>
#
# And the model may say something to the user on the commentary channel
# before a call, with no recipient:
#
#   <|channel|>commentary<|message|>Let me check.<|end|><|start|>assistant<|channel|>commentary to=…
#
# mlx_lm 0.31.3 knows none of this, so the whole thing came back as the
# reply's text, thinking and markup included. It does know thinking and
# tool calls in general, as a start marker and an end marker each, read
# off the tokenizer, and its own state machine drops the marker text. So
# the format is taught to the tokenizer in those terms: the analysis
# header opens thinking and `<|end|>` closes it; `to=functions`, with or
# without the commentary header in front of it, opens a tool call that
# the `<|call|>` stop token ends. Everything else the format writes
# between messages (`<|start|>assistant`, the final header, a preamble's
# header, an `<|end|>` outside thinking) is dropped by the state machine
# and the text between stays. That leaves a preamble in the reply's text,
# which is what it is.
#
# A call that ends on a stop token needs one more thing. mlx_lm's handler
# files a tool call when the state after it is "normal", and a stop has
# no state, so a reply that stops inside a call would be filed as nothing.
# `Generator.generate` below hands the handler the stop with the state
# "normal", and the call is filed.

HARMONY_ANALYSIS = "<|channel|>analysis<|message|>"
HARMONY_CLOSE = "<|end|>"
HARMONY_TURN = "<|start|>assistant"
HARMONY_FINAL = "<|channel|>final<|message|>"
HARMONY_PREAMBLE = "<|channel|>commentary<|message|>"
# All end at `functions`: the `.` after it is one token with what follows.
# The format puts function calls in the commentary channel, but gpt-oss
# also calls them straight from the analysis channel, in the middle of
# thinking, and a call is a call whichever channel it came from.
HARMONY_TOOL = "<|channel|>commentary to=functions"
HARMONY_ANALYSIS_TOOL = "<|channel|>analysis to=functions"
HARMONY_RECIPIENT = " to=functions"
HARMONY_TOKENS = ("<|channel|>", "<|message|>", "<|start|>", "<|end|>", "<|call|>")

# The text of a tool call after its opening: the function's name, then
# the commentary header when the recipient came in the role header, then
# an optional content type with or without `<|constrain|>`, then the
# arguments after `<|message|>`.
HARMONY_TOOL_CALL = re.compile(
    r"\s*\.?(?P<name>[\w\-]+)\s*(?:<\|channel\|>commentary)?\s*(?:<\|constrain\|>)?[\w/+.\-]*\s*<\|message\|>(?P<args>.*)$",
    re.DOTALL,
)


def speaks_harmony(tokenizer):
    """Whether the tokenizer has every token the format is spelled with."""
    vocab = tokenizer.get_vocab()
    return all(token in vocab for token in HARMONY_TOKENS)


def parse_harmony_tool_call(text, tools):
    """The tool call in the text between `to=functions` and the `<|call|>`
    that ended the reply, in the shape mlx_lm's parsers return. Arguments
    that are not JSON are handed over as they were written, under `raw`,
    rather than failing the request."""
    match = HARMONY_TOOL_CALL.match(text)
    if match is None:
        raise ValueError(f"Not a Harmony tool call: {text[:80]!r}")
    args_text = match.group("args").strip()
    try:
        arguments = json.loads(args_text) if args_text else {}
    except json.JSONDecodeError:
        arguments = {"raw": args_text}
    return {"name": match.group("name"), "arguments": arguments}


def teach_harmony(tokenizer):
    """Give a gpt-oss tokenizer its thinking and tool-call markers, once.
    The wrapper keeps every underscore attribute on itself, so these land
    where its properties read them."""
    if getattr(tokenizer, "_harmony", False) or not speaks_harmony(tokenizer):
        return
    encode = lambda text: tuple(tokenizer.encode(text, add_special_tokens=False))
    tokenizer._think_start = HARMONY_ANALYSIS
    tokenizer._think_end = HARMONY_CLOSE
    tokenizer._think_start_tokens = encode(HARMONY_ANALYSIS)
    tokenizer._think_end_tokens = encode(HARMONY_CLOSE)
    tokenizer._tool_parser = parse_harmony_tool_call
    tokenizer._tool_call_start = HARMONY_TOOL
    tokenizer._tool_call_start_tokens = encode(HARMONY_TOOL)
    tokenizer._tool_call_end = None
    tokenizer._tool_call_end_tokens = None
    # The other ways a call can open, and the markers dropped between
    # messages. The final header is also the answer's header for the
    # reply watcher: on its own at the reply's start, after the turn
    # marker once thinking has closed.
    tokenizer._harmony_tool_openers = [encode(HARMONY_TOOL), encode(HARMONY_ANALYSIS_TOOL), encode(HARMONY_RECIPIENT)]
    tokenizer._harmony_drops = {
        encode(HARMONY_TURN): HARMONY_TURN,
        encode(HARMONY_FINAL): HARMONY_FINAL,
        encode(HARMONY_PREAMBLE): HARMONY_PREAMBLE,
        encode(HARMONY_CLOSE): HARMONY_CLOSE,
    }
    tokenizer._harmony_headers = [encode(HARMONY_TURN + HARMONY_FINAL), encode(HARMONY_FINAL)]
    tokenizer._harmony = True
    logging.info("The model speaks Harmony; its channels will be split.")


def harmony_template_kwargs(kwargs):
    """gpt-oss's template has no thinking switch, only `reasoning_effort`.
    A request that turned thinking off gets the lowest effort, unless it
    named one."""
    kwargs = dict(kwargs or {})
    if kwargs.get("enable_thinking") is False and "reasoning_effort" not in kwargs:
        kwargs["reasoning_effort"] = "low"
    return kwargs


# ---------------------------------------------------------------------------
# The seams into mlx_lm.server.
# ---------------------------------------------------------------------------


class Generator(server.ResponseGenerator):
    """The one place that sees the request, the tokenizer, and the prompt's
    thinking state together. It hangs the watcher on the request's
    arguments, where `make_logits_processors` picks it up. It also hands
    each handler a token stream that watches the handler's socket."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ll_tokenizers = {}
        self.token_masks = {}
        self.pieces = {}

    def ll_tokenizer_for(self, tokenizer, vocab_width):
        # Wrapping a tokenizer takes about a second, so each is kept, keyed
        # on the model it belongs to.
        key = (self.model_provider.model_key, vocab_width)
        if key not in self.ll_tokenizers:
            self.ll_tokenizers[key] = from_tokenizer(
                tokenizer._tokenizer,
                n_vocab=vocab_width,
                eos_token=list(tokenizer.eos_token_ids),
            )
        return self.ll_tokenizers[key]

    def token_masks_for(self, tokenizer, vocab_width):
        # Kept for the same reason, keyed the same way.
        key = (self.model_provider.model_key, vocab_width)
        if key not in self.token_masks:
            self.token_masks[key] = token_masks(tokenizer._tokenizer, vocab_width)
        return self.token_masks[key]

    def piece_for(self, tokenizer, token_id):
        """The text of one token on its own. Rough at the edges of a
        multi-byte character, which is fine for scanning a reply for
        phrases and sentences, and cheap enough to keep per token."""
        key = (self.model_provider.model_key, token_id)
        piece = self.pieces.get(key)
        if piece is None:
            piece = tokenizer.decode([token_id])
            self.pieces[key] = piece
        return piece

    def vocab_width(self):
        """The main model's output width, from its config."""
        args = getattr(self.model_provider.model, "args", None)
        return getattr(args, "vocab_size", 0) or 0

    def _make_state_machine(self, model_key, tokenizer, stop_words, initial_state="normal"):
        """mlx_lm's, rebuilt for a Harmony tokenizer: a tool call can open
        three ways, and the markers between messages are dropped, which is a
        rule from the normal state back to itself. mlx_lm blanks the text of
        any matched sequence, so that is all a drop takes."""
        teach_harmony(tokenizer)
        if not getattr(tokenizer, "_harmony", False):
            return super()._make_state_machine(model_key, tokenizer, stop_words, initial_state)
        cache_key = ("harmony", model_key, tuple(stop_words), initial_state)
        cached = self._state_machine_cache.get(cache_key)
        if cached is not None:
            return cached
        sequences = {}
        stops = []
        for token in tokenizer.eos_token_ids:
            sequences[(token,)] = tokenizer.convert_ids_to_tokens(token)
            stops.append(((token,), None))
        for word in stop_words:
            tokens = tuple(tokenizer.encode(word, add_special_tokens=False))
            sequences[tokens] = word
            stops.append((tokens, None))
        think_start = tokenizer.think_start_tokens
        think_end = tokenizer.think_end_tokens
        sequences[think_start] = tokenizer.think_start
        sequences[think_end] = tokenizer.think_end
        normal = [(think_start, "reasoning")]
        for opener in tokenizer._harmony_tool_openers:
            sequences[opener] = tokenizer.decode(list(opener))
            normal.append((opener, "tool"))
        for dropped, text in tokenizer._harmony_drops.items():
            sequences[dropped] = text
            normal.append((dropped, "normal"))
        transitions = {
            "normal": normal + stops,
            "reasoning": [(think_end, "normal")] + stops,
            "tool": list(stops),
        }
        machine = server.SequenceStateMachine(transitions, initial=initial_state)
        if len(self._state_machine_cache) > 100:
            self._state_machine_cache.clear()
        self._state_machine_cache[cache_key] = (machine, sequences)
        return machine, sequences

    def _tokenize(self, tokenizer, request, args):
        teach_harmony(tokenizer)
        if getattr(tokenizer, "_harmony", False):
            args.chat_template_kwargs = harmony_template_kwargs(args.chat_template_kwargs)
        result = super()._tokenize(tokenizer, request, args)
        initial_state = result[3]
        grammar = getattr(request, "grammar", None)
        limits = getattr(request, "limits", None)
        if limits is None:
            limits = limits_for({}, args.max_tokens, SERVER_LIMITS)
        if grammar is not None:
            logging.info(f"Constraining the reply to a schema, from state {initial_state!r}")
        args.watcher = Watcher(
            tokenizer,
            lambda width: self.ll_tokenizer_for(tokenizer, width),
            lambda width: self.token_masks_for(tokenizer, width),
            lambda token_id: self.piece_for(tokenizer, token_id),
            grammar,
            initial_state,
            limits,
            self.vocab_width(),
        )
        return result

    def _serve_single(self, request):
        """mlx_lm's path for a request it cannot batch: every request when a
        draft model is loaded, and any request with a seed. mlx_lm holds the
        prompt cache to `--prompt-cache-bytes` only on the batched path, so
        this one trims it the same way: before the reply, so the cache is
        under the cap while the reply runs, and after, for what the reply
        added."""
        limit = self.model_provider.cli_args.prompt_cache_bytes
        if limit is not None:
            self.prompt_cache.trim_to(n_bytes=limit)
        super()._serve_single(request)
        if limit is not None:
            self.prompt_cache.trim_to(n_bytes=limit)

    def check_generating(self):
        """Every reply comes from the one generation thread. If it has died,
        from an error it did not catch such as the GPU running out of memory,
        every request would wait for ever. Fail the request instead, so the
        caller sees an error and the server can be restarted."""
        if not self._generation_thread.is_alive():
            raise RuntimeError("The server's generation thread has stopped; restart the server.")

    def generate(self, request, args, progress_callback=None):
        """mlx_lm's `generate`, with two changes. While the handler waits for
        the next token it looks at its socket every half second, and stops
        the reply when the client has closed it. mlx_lm's version blocks on
        the queue until the reply is complete, so it can only find out the
        client is gone when it writes the finished reply. And a reply the
        watcher cut short ends with finish_reason "length", not "stop".

        The generation thread stops a reply by dropping it from the batch
        on its next step, once `ctx.stop()` has been called; the tokens it
        already queued are left unread."""
        connection = getattr(request, "connection", None)
        response_queue = Queue()
        self.requests.put((response_queue, request, args))

        while True:
            try:
                ctx = response_queue.get(timeout=CLIENT_POLL_SECONDS)
                break
            except QueueEmpty:
                self.check_generating()
        if isinstance(ctx, Exception):
            raise ctx

        def tokens():
            # The socket is looked at on a clock, not only when the queue
            # runs dry: a reply that streams a token every few milliseconds
            # never leaves the queue empty for long.
            last_look = time.monotonic()
            # The state of the last token handed over, so a stop inside a
            # tool call can be preceded by the step that files the call.
            last_state = None
            while True:
                try:
                    response = response_queue.get(timeout=CLIENT_POLL_SECONDS)
                except QueueEmpty:
                    response = WAITING
                    self.check_generating()
                if connection is not None and time.monotonic() - last_look >= CLIENT_POLL_SECONDS:
                    last_look = time.monotonic()
                    if client_gone(connection):
                        logging.info("The client went away; stopping its reply.")
                        ctx.stop()
                        return
                if response is WAITING:
                    continue
                if response is None:
                    return
                if isinstance(response, Exception):
                    raise response
                if isinstance(response, tuple):
                    if progress_callback is not None:
                        progress_callback(*response)
                    continue
                watcher = getattr(args, "watcher", None)
                if response.finish_reason is not None and watcher is not None and watcher.cut:
                    response = replace(response, finish_reason="length")
                response = file_tool_call_on_stop(last_state, response)
                last_state = response.state
                yield response

        return ctx, server._process_control_tokens(ctx, tokens())


def file_tool_call_on_stop(last_state, response):
    """A stop met inside a tool call, handed over with the state "normal"
    so mlx_lm's handler files the call. A tool call that ends on a stop
    token (Harmony's `<|call|>`) stops the reply while the state machine
    is still in the tool state, and the handler only files a call when a
    "normal" token follows it. The stop's own text is already blank, and
    the handler reads the finish reason from the token itself, so nothing
    else changes; in particular no token is added to the reply's count.
    Any other token is returned as it is."""
    if response.finish_reason is None or response.state is not None or last_state != "tool":
        return response
    return replace(response, state="normal")


# ---------------------------------------------------------------------------
# A Metal buffer leak in mlx_lm 0.31.3's batched decoding.
# ---------------------------------------------------------------------------
#
# Models with layers that keep state other than keys and values (mlx_lm's
# `ArraysCache`: Qwen3.5 and its descendants, Qwen3-Next, Mamba, LFM2,
# Kimi Linear and others) advance that cache each step with a lazy
# subtraction on its padding and length counters, so each step's counter
# keeps the graph that made it, and every Metal buffer in that chain
# stays alive: one per such layer per token. Metal allows 499,000 live
# buffers per process, which a reply reaches after about 10,000 generated
# tokens; the generation thread then dies with
# `[metal::malloc] Resource limit (499000) exceeded` and the server answers
# every later request with a 404 (ml-explore/mlx-lm #1641, #1332, #1672;
# fixed upstream in #1642, unreleased as of 0.31.3). Evaluating the
# counters, and the state beside them, after each step collapses the
# chain. They are already computed by the step, so this costs a
# synchronisation, not work.


def cache_arrays(caches):
    """Every array a cache keeps: its state, however nested, and the
    padding and length counters an `ArraysCache` advances lazily."""
    found = []

    def collect(value):
        if isinstance(value, mx.array):
            found.append(value)
        elif isinstance(value, (list, tuple)):
            for item in value:
                collect(item)

    for cache in caches:
        collect(getattr(cache, "state", None))
        collect(getattr(cache, "left_padding", None))
        collect(getattr(cache, "lengths", None))
    return found


# Remove this shim once a released mlx_lm carries the upstream fix (#1642);
# the extra per-step eval is redundant after that. Unwrap before capturing so
# that if this module is ever re-executed, `original_batch_step` is the true
# original and `settled_step` never wraps itself into infinite recursion.
_installed = GenerationBatch._step
original_batch_step = getattr(_installed, "__cache_leak_original__", _installed)


def settled_step(self):
    result = original_batch_step(self)
    arrays = cache_arrays(self.prompt_cache)
    if arrays:
        mx.eval(arrays)
    return result


settled_step.__cache_leak_original__ = original_batch_step
GenerationBatch._step = settled_step


original_make_logits_processors = server._make_logits_processors


def lenient_param_value(convert):
    """A tool-call parameter that does not read as its declared type is
    handed over as the text the model wrote, instead of failing the whole
    request. mlx_lm 0.31.3's Qwen parser evaluates a parameter whose type
    it does not know as a Python literal, which a plain word is not, and
    reads an object parameter as JSON, which fails when the model wrote
    anything after the closing brace. A number parameter is read as
    int(float(value)), so a value too large to hold, such as 1e400,
    overflows to infinity and raises OverflowError. Any of these raised
    out of the parser, and the server answered 502 to the agent's tool
    call. The agent's tools take strings and check them, so the text is
    the right fallback. Filed against mlx-lm in the benchmark's
    upstream-bug-reports.md."""

    def converted(param_value, param_name, param_config):
        try:
            return convert(param_value, param_name, param_config)
        except PARSER_ERRORS:
            return param_value

    return converted


# Returned by `read_as_member` when the text is not a value of that type.
NOT_THAT_TYPE = object()

STRING_TYPES = ("string", "str", "text")

# What the parser returns for each JSON Schema type a union can name. It
# reads an array written as a Python tuple, so a tuple counts as one.
MEMBER_VALUES = {
    "integer": int,
    "number": (int, float),
    "boolean": bool,
    "object": dict,
    "array": (list, tuple),
}

PARSER_ERRORS = (ValueError, SyntaxError, TypeError, OverflowError)


def union_member_types(param):
    """The type names a parameter's schema allows when it is a union, such
    as ["number", "null"] for `number | null`, or an empty list when the
    schema names one type or none. zod writes a union as `anyOf`; JSON
    Schema also allows `oneOf` and a list under `type`."""
    if not isinstance(param, dict):
        return []
    declared = param.get("type")
    if isinstance(declared, list):
        return [str(name).lower() for name in declared]
    if declared is not None:
        return []
    members = param.get("anyOf") or param.get("oneOf") or []
    names = [member.get("type") for member in members if isinstance(member, dict)]
    return [str(name).lower() for name in names if isinstance(name, str)]


def read_as_member(convert, text, type_name):
    """The text as a value of one member type, or NOT_THAT_TYPE. The parser
    does the reading, as if the parameter had that one type, so a union
    member is read exactly as a parameter of that type is."""
    wanted = MEMBER_VALUES.get(type_name)
    if wanted is None:
        return NOT_THAT_TYPE
    # The parser reads every text but "true" as false, so it cannot say
    # whether the text is a boolean at all.
    if type_name == "boolean" and text.lower() not in ("true", "false"):
        return NOT_THAT_TYPE
    try:
        value = convert(text, "member", {"member": {"type": type_name}})
    except PARSER_ERRORS:
        return NOT_THAT_TYPE
    return value if isinstance(value, wanted) else NOT_THAT_TYPE


def union_param_value(convert):
    """Reads a tool-call parameter whose type is a union as the first
    member the text fits. mlx_lm 0.31.3's Qwen parser looks for one `type`
    on the parameter's schema. A union has none, so the parser hands a
    number sent to a `number | null` parameter over as text.

    The word "null" is nothing, as it is for every parameter the parser
    reads. A union that allows a string keeps any other text, since any
    text is a valid string. Text that fits no member also stays text, so
    the tool can tell the model what was wrong with it."""

    def converted(param_value, param_name, param_config):
        members = union_member_types(param_config.get(param_name))
        if not members:
            return convert(param_value, param_name, param_config)
        # With no schema to go on, the parser applies only its rule for
        # the word "null".
        if convert(param_value, param_name, {}) is None:
            return None
        if any(name in STRING_TYPES for name in members):
            return param_value
        for name in members:
            value = read_as_member(convert, param_value, name)
            if value is not NOT_THAT_TYPE:
                return value
        return param_value

    return converted


def make_tool_parsers_lenient():
    """Wrap the parser that has been seen to fail. Other parsers are left
    alone until one is seen to."""
    try:
        from mlx_lm.tool_parsers import qwen3_coder
    except ImportError:
        return
    qwen3_coder._convert_param_value = lenient_param_value(
        union_param_value(qwen3_coder._convert_param_value)
    )


def keep_logits(tokens, logits):
    return logits


def make_logits_processors(args):
    """The request's own processors, plus its watcher. A request with
    nothing to apply still gets one processor that changes nothing: mlx_lm
    0.31.3's batch generator trims its list of processors only when some
    entry is truthy, so a request with an empty list would leave a stale
    entry behind when it finished, and the next request's watcher would
    sit at the wrong index and never run. Fixed upstream in mlx-lm PR
    #1772, after 0.31.3; drop this once the pin moves past it."""
    processors = list(original_make_logits_processors(args))
    watcher = getattr(args, "watcher", None)
    if watcher is not None:
        processors.append(watcher)
    if not processors:
        processors.append(keep_logits)
    return processors


def one_system_message(messages):
    """The messages with every system message folded into one at the front.
    A chat template takes one system message and wants it first; the
    Qwen3.8 template refuses the request otherwise. An OpenAI-style client
    can send several, and a client speaking the Anthropic API through a
    translating proxy does: Claude Code's system prompt arrives as one
    system message per block. Their texts are joined with blank lines, in
    order, and the rest of the conversation keeps its order."""
    if not isinstance(messages, list):
        return messages
    system = [m for m in messages if isinstance(m, dict) and m.get("role") == "system"]
    if len(system) <= 1 and (not system or messages[0] is system[0]):
        return messages
    texts = []
    for message in system:
        content = message.get("content")
        if isinstance(content, list):
            content = "".join(part.get("text", "") for part in content if isinstance(part, dict))
        if content:
            texts.append(str(content))
    rest = [m for m in messages if not (isinstance(m, dict) and m.get("role") == "system")]
    return [{"role": "system", "content": "\n\n".join(texts)}] + rest


# What the OpenAI-style `reasoning_effort` field means to a template that
# knows three levels. Effort names outside the three are the nearest one.
EFFORT_LEVELS = {
    "none": "low", "minimal": "low", "low": "low",
    "medium": "medium",
    "high": "high", "xhigh": "high", "max": "high",
}


def fold_reasoning_effort(body):
    """The body with a top-level `reasoning_effort` carried into its
    `chat_template_kwargs`, where gpt-oss's template reads it. Clients
    send the effort as its own field, the way the OpenAI API takes it;
    the template only sees template kwargs. An explicit kwarg wins, and
    a level the template does not know becomes the nearest one it does.
    Other templates ignore a kwarg they have no use for."""
    if not isinstance(body, dict):
        return body
    effort = EFFORT_LEVELS.get(str(body.get("reasoning_effort") or "").lower())
    if effort is None:
        return body
    kwargs = dict(body.get("chat_template_kwargs") or {})
    kwargs.setdefault("reasoning_effort", effort)
    return {**body, "chat_template_kwargs": kwargs}


class Handler(server.APIHandler):
    def handle_chat_completions(self):
        if isinstance(self.body, dict):
            self.body["messages"] = one_system_message(self.body.get("messages"))
            self.body = fold_reasoning_effort(self.body)
        request = super().handle_chat_completions()
        request.grammar = None
        # The socket the reply goes to, so `Generator.generate` can tell
        # when the client has closed it.
        request.connection = self.connection
        try:
            request.limits = limits_for(self.body, self.max_tokens, SERVER_LIMITS)
            if schema_can_be_enforced(request.tools, self.body.get("tool_choice")):
                request.grammar = grammar_for(self.body.get("response_format"))
            elif self.body.get("response_format") is not None:
                logging.warning("response_format is not enforced on a request with tools.")
        except ValueError as problem:
            self._set_completion_headers(400)
            self.end_headers()
            self.wfile.write(json.dumps({"error": str(problem)}).encode())
            raise RequestRefused()
        return request

    def do_POST(self):
        try:
            super().do_POST()
        except RequestRefused:
            return
        except (BrokenPipeError, ConnectionResetError):
            # A stopped reply is still written out, to a socket the client
            # has closed. Without this, the server prints a traceback for it.
            logging.info("The client went away before its reply was written.")


class RequestRefused(Exception):
    """The handler already wrote a 400; nothing else should be written."""


def untrimmable_cache(cache):
    """Why speculative decoding cannot run on a model with this attention
    cache, or None when it can. The main model checks a whole draft in one
    pass and has to take back the guesses it rejects, so every layer's cache
    must be able to give tokens back. The standard attention cache can;
    the recurrent and hybrid ones (Qwen3.5, Qwen3-Next, Mamba) cannot."""
    if can_trim_prompt_cache(cache):
        return None
    kinds = ", ".join(sorted({type(c).__name__ for c in cache}))
    return f"its attention cache ({kinds}) cannot give tokens back"


def refuse_draft_if_unusable(model_provider):
    """mlx_lm only finds out that a model cannot draft when the first
    request fails, after the server has answered its readiness check. This
    finds out before the server opens, so `agency local serve --draft` fails
    with the reason instead of failing every request."""
    if model_provider.cli_args.draft_model is None:
        return
    model_provider.load_default()
    if model_provider.draft_model is None:
        return
    problem = untrimmable_cache(server.make_prompt_cache(model_provider.model))
    if problem is None:
        return
    logging.error(
        f"{model_provider.cli_args.model} cannot use a draft model: {problem}. "
        f"Serve it without --draft."
    )
    sys.exit(2)


def run(host, port, model_provider):
    refuse_draft_if_unusable(model_provider)
    prompt_cache = server.LRUPromptCache(model_provider.cli_args.prompt_cache_size)
    generator = Generator(model_provider, prompt_cache)
    server._run_http_server(host, port, generator, handler_class=Handler)


def split_own_flags(argv):
    """The limit flags, which are this script's own, and the rest of the
    command line, which is mlx_lm.server's."""
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--reasoning-budget", type=int, default=SERVER_LIMITS.reasoning_budget)
    parser.add_argument("--hedge-limit", type=int, default=SERVER_LIMITS.hedge_limit)
    parser.add_argument("--repeat-limit", type=int, default=SERVER_LIMITS.repeat_limit)
    parser.add_argument("--limit-answers", action="store_true", default=SERVER_LIMITS.limit_answers)
    own, rest = parser.parse_known_args(argv)
    limits = Limits(
        reasoning_budget=own.reasoning_budget,
        hedge_limit=own.hedge_limit,
        repeat_limit=own.repeat_limit,
        limit_answers=own.limit_answers,
    )
    return limits, rest


sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from localServerCommon import exit_when_parent_goes  # noqa: E402

if __name__ == "__main__":
    exit_when_parent_goes()
    SERVER_LIMITS, rest = split_own_flags(sys.argv[1:])
    sys.argv = [sys.argv[0]] + rest
    server.run = run
    server._make_logits_processors = make_logits_processors
    make_tool_parsers_lenient()
    server.main()