fix(studio): count an MLX prompt with the resident model's own tokenizer

Studio's chat header prices a conversation against the model's context window, and the recount behind it had a llama.cpp branch only. For an MLX-served model the endpoint answered "No GGUF model loaded", so the bar showed a dash for the whole conversation until the first reply reported its own usage -- the case unslothai#8882 made visible by rendering the bar as soon as the window is known.

The endpoint now asks MLX after finding no GGUF backend, below the refusals that are about the request rather than the backend and above llama.cpp's own render concerns.

A count is only worth showing if it renders what the completion for the same request would render, so this answers the questions that completion answers rather than the ones the request makes easy. Which of its two paths claims the request: the tool loop, or the relay that a client catalog or replayed tool history goes to. What each then renders: the loop's action nudge and its strip of stale call markup, the relay's rebuild that keeps structured tool_calls through templating, the launcher's tools-on default where a request stated no intent of its own and its withdrawal where the request did, the catalog a named template's tool_use branch advertises, and a zero tool budget suppressing the loop entirely. Where a helper carries one of those decisions it is called rather than restated.

Three shapes are refused rather than priced short: a pending turn whose retrieval would reach the document store, since only running the search would say what it adds; an image, which /apply-template swaps for a short marker; and an empty prompt, whose bare generation marker reads as a started conversation.

A vision model is counted rather than skipped. It serves text turns through the processor render its generation uses, so counting shares that render, and asks mlx_vlm which special markers tokenization should add -- a per-model answer, and on the releases from before mlx_vlm exported it, the rule those releases inline.

Both store questions on this path -- the enabled MCP servers and whether retrieval can run -- open SQLite, and the second can load an extension and create the schema, so they run on the worker thread and only where their answer decides something.

Author: Lyxot

studio/backend/core/inference/mlx_inference.py

@@ -1127,6 +1127,15 @@ def _processor(tokens, logits):
     return _processor
 
 
+# The families the mlx_vlm releases from before should_add_special_tokens existed inlined into
+# their generation path: their chat template emits the special markers, so tokenization must not
+# add them again. Deliberately not the current helper's list, which also carries laguna -- laguna
+# models run on 0.6.0 onward but only 0.6.9 stopped tokenizing them with the markers, and 0.6.0
+# through 0.6.8 are what this stands in for. A family that arrived with its list entry is safe to
+# keep here, since a release without the entry cannot load the model either.
+_VLM_INLINE_SPECIAL_TOKEN_FAMILIES = ("gemma3", "gemma3n", "gemma4", "gemma4_unified")
+
+
 class MLXInferenceBackend:
     def __init__(self):
         self.models = {}
@@ -1724,13 +1733,37 @@ def count_chat_tokens(
         """
         if self._model is None:
             raise RuntimeError("No model loaded")
-        if self._is_vlm:
-            # A vision model renders through its processor, which recovers from template
-            # failures by rewriting the conversation; the text renderer would not.
-            raise RuntimeError("Counting is not supported for vision models")
-
         full_messages = self._with_system_prompt(messages, system_prompt)
 
+        if self._is_vlm:
+            # Through the processor, which is what a vision generation renders with; the
+            # text renderer would not recover the template failures it recovers from.
+            # images=None: an image anywhere in the conversation makes the structured-item
+            # check raise, and the caller declines rather than pricing a prompt without it.
+            prompt, _ = self._render_vlm_prompt(
+                full_messages,
+                None,
+                tools = tools,
+                enable_thinking = enable_thinking,
+                reasoning_effort = reasoning_effort,
+                preserve_thinking = preserve_thinking,
+            )
+            # Whether the markers belong to the template or to tokenization is a per-model
+            # answer mlx_vlm makes for every generation; ask it rather than guess, or the
+            # count is off by whatever the generation's own choice would have added.
+            _model_type = getattr(getattr(self._model, "config", None), "model_type", None)
+            try:
+                from mlx_vlm.utils import should_add_special_tokens
+                add_special = should_add_special_tokens(_model_type, self._processor)
+            except Exception:
+                # The rule those releases inline, which Studio's runtime gate still accepts.
+                add_special = (
+                    getattr(self._processor, "chat_template", None) is None
+                    if _model_type in _VLM_INLINE_SPECIAL_TOKEN_FAMILIES
+                    else True
+                )
+            return len(self._tokenizer.encode(prompt, add_special_tokens = add_special))
+
         render_result = self._render_text_prompt(
             full_messages,
             tools = tools,
@@ -1988,28 +2021,22 @@ def _generate_text(
                 normalized_output += tail
                 yield normalized_output
 
-    def _generate_vlm(
+    def _render_vlm_prompt(
         self,
         messages,
-        image,
-        temperature,
-        top_p,
-        top_k,
-        min_p,
-        max_new_tokens,
-        repetition_penalty,
-        cancel_event,
+        images,
         *,
         tools = None,
         enable_thinking = None,
         reasoning_effort = None,
         preserve_thinking = None,
         continue_final_message = False,
-        presence_penalty = 0.0,
-        _adapter_state = None,
     ):
-        from mlx_vlm import stream_generate as vlm_stream
+        """Render the prompt a vision generation sends, and the target that rendered it.
 
+        Shared with counting, as _render_text_prompt is for text models, so a count cannot
+        price a prompt the model never sees. A text-only conversation passes images=None.
+        """
         from core.inference.chat_template_helpers import (
             apply_chat_template_for_generation,
             chat_render_target,
@@ -2021,8 +2048,6 @@ def _generate_vlm(
         # to authorize against the same template this line selects (#7066).
         chat_target = chat_render_target(self._processor)
 
-        # mlx_vlm's stream_generate handles pixel_values (None for text-only)
-        images = [image] if image is not None else None
         attached_images = 0 if images is None else len(images)
         structured_images = sum(
             _count_vlm_images(message.get("content"))
@@ -2104,6 +2129,40 @@ def _generate_vlm(
             prompt = recovered_prompt
         elif prompt_issue:
             raise RuntimeError(f"VLM chat template returned {prompt_issue}.") from prompt_error
+        return prompt, chat_target
+
+    def _generate_vlm(
+        self,
+        messages,
+        image,
+        temperature,
+        top_p,
+        top_k,
+        min_p,
+        max_new_tokens,
+        repetition_penalty,
+        cancel_event,
+        *,
+        tools = None,
+        enable_thinking = None,
+        reasoning_effort = None,
+        preserve_thinking = None,
+        continue_final_message = False,
+        presence_penalty = 0.0,
+        _adapter_state = None,
+    ):
+        from mlx_vlm import stream_generate as vlm_stream
+
+        images = [image] if image is not None else None
+        prompt, chat_target = self._render_vlm_prompt(
+            messages,
+            images,
+            tools = tools,
+            enable_thinking = enable_thinking,
+            reasoning_effort = reasoning_effort,
+            preserve_thinking = preserve_thinking,
+            continue_final_message = continue_final_message,
+        )
 
         from core.inference.chat_template_helpers import detect_think_prefill
 

studio/backend/models/inference.py

@@ -2044,22 +2044,41 @@ class ChatCountTokensRequest(ReasoningControlsRequest):
     )
     permission_mode: Optional[str] = Field(
         None,
-        description = "[x-unsloth] Permission level the completion would send. Only 'full' changes "
-        "the prompt: it swaps the python/terminal descriptions for the unsandboxed pair and adds a "
-        "sentence to the tool nudge, so a count that omits it prices a prompt the completion will "
-        "not send.",
+        description = "[x-unsloth] Permission level the completion would send. 'full' swaps the "
+        "python/terminal descriptions for the unsandboxed pair and adds a sentence to the tool "
+        "nudge; 'ask' holds the tool loop's first-pass retrieval behind its confirmation gate, so "
+        "a pending turn under a retrieval scope is countable. A count that omits this prices a "
+        "prompt the completion will not send, or declines one it could have priced.",
     )
     bypass_permissions: Optional[bool] = Field(
         None,
         description = "[x-unsloth] Equivalent of permission_mode='full'. Declared explicitly (not "
         "left to extra='allow') so an omitted flag reads as None instead of raising AttributeError.",
     )
 
+    confirm_tool_calls: Optional[bool] = Field(
+        None,
+        description = "[x-unsloth] Whether the completion's tool loop would gate each call. "
+        "Declared so it is typed and folded into permission_mode as it is on a completion.",
+    )
+    max_tool_calls_per_message: Optional[int] = Field(
+        None,
+        ge = 0,
+        description = "[x-unsloth] Tool-call budget the completion would send. Zero suppresses the "
+        "tool loop, so a count that never sees it prices a catalog the relay does not render.",
+    )
+
     @field_validator("permission_mode", mode = "before")
     @classmethod
     def _coerce_permission_mode(cls, value: Any) -> Any:
         return _normalize_permission_mode(value)
 
+    # The very function the completion request runs, not a copy: a count renders replayed
+    # tool history through the same templates, which read the id off the result message.
+    _resolve_missing_tool_call_ids = model_validator(mode = "after")(
+        ChatCompletionRequest._resolve_missing_tool_call_ids
+    )
+
     @model_validator(mode = "after")
     def _fold_full_permission_into_bypass(self) -> "ChatCountTokensRequest":
         """Mirrors ChatCompletionRequest: the prompt builders read only the
@@ -2068,6 +2087,11 @@ def _fold_full_permission_into_bypass(self) -> "ChatCountTokensRequest":
             self.bypass_permissions = True
         elif self.bypass_permissions:
             self.permission_mode = "full"
+        elif self.permission_mode is None and self.confirm_tool_calls is True:
+            # The same reading a completion gives it: gating every call is the
+            # pre-permission-mode way of asking for "ask", and the loop's retrieval
+            # gate turns on that. No provider clause -- this endpoint is local only.
+            self.permission_mode = "ask"
         return self