Studio: add reasoning budget controls

studio/backend/core/inference/llama_cpp.py

@@ -52,6 +52,8 @@
     parse_gpu_layers_override,
     parse_split_mode_override,
     resolve_requested_ctx,
+    resolve_reasoning_budget,
+    resolve_reasoning_budget_message,
     strip_shadowing_flags,
     strip_split_mode_only,
 )
@@ -2412,6 +2414,8 @@ def __init__(self):
         self._chat_template_override: Optional[str] = None
         self._supports_reasoning: bool = False
         self._reasoning_always_on: bool = False
+        self._reasoning_budget: int = -1
+        self._reasoning_budget_message: str = ""
         self._reasoning_style: str = "enable_thinking"
         self._reasoning_effort_levels: list = []
         self._supports_preserve_thinking: bool = False
@@ -2625,6 +2629,14 @@ def requested_n_ctx(self) -> int:
         0 means Auto. Used by the route to detect Auto-vs-explicit flips."""
         return self._requested_n_ctx
 
+    @property
+    def reasoning_budget(self) -> int:
+        return self._reasoning_budget
+
+    @property
+    def reasoning_budget_message(self) -> str:
+        return self._reasoning_budget_message
+
     @property
     def extra_args_source(self) -> Optional[tuple[str, Optional[str]]]:
         """(model_identifier, hf_variant) the stored extra_args came from.
@@ -5920,6 +5932,8 @@ def _start_diffusion_server(
         self._is_audio = False  # clear any prior TTS/audio model's routing flag
         self._model_identifier = model_identifier
         self._cache_type_kv = None
+        self._reasoning_budget = -1
+        self._reasoning_budget_message = ""
         self._swa_full = False
         self._kv_cache_unified = False
         self._n_ubatch = self._DEFAULT_N_UBATCH
@@ -7156,6 +7170,8 @@ def load_model(
         cache_type_kv: Optional[str] = None,
         speculative_type: Optional[str] = None,
         spec_draft_n_max: Optional[int] = None,
+        reasoning_budget: int = -1,
+        reasoning_budget_message: str = "",
         tensor_parallel: bool = False,
         gpu_memory_mode: Literal["auto", "manual"] = "auto",
         gpu_layers: int = -1,
@@ -7196,6 +7212,8 @@ def load_model(
             "cache_type_kv": cache_type_kv,
             "speculative_type": speculative_type,
             "spec_draft_n_max": spec_draft_n_max,
+            "reasoning_budget": reasoning_budget,
+            "reasoning_budget_message": reasoning_budget_message,
             "tensor_parallel": tensor_parallel,
             # GPU-memory placement: replayed on respawn so a server SIGKILL'd by
             # GPU/RAM pressure reloads onto the same devices with the same
@@ -7231,6 +7249,8 @@ def load_model(
                 cache_type_kv = cache_type_kv,
                 speculative_type = speculative_type,
                 spec_draft_n_max = spec_draft_n_max,
+                reasoning_budget = reasoning_budget,
+                reasoning_budget_message = reasoning_budget_message,
                 tensor_parallel = tensor_parallel,
                 gpu_memory_mode = gpu_memory_mode,
                 gpu_layers = gpu_layers,
@@ -9006,6 +9026,16 @@ def _restore_after_tensor_downgrade():
                     )
                     logger.info(f"Reasoning model: {reasoning_kw} by default")
 
+                reasoning_budget = resolve_reasoning_budget(extra_args, reasoning_budget)
+                reasoning_budget_message = resolve_reasoning_budget_message(
+                    extra_args, reasoning_budget_message
+                )
+                cmd.extend(["--reasoning-budget", str(reasoning_budget)])
+                if reasoning_budget_message:
+                    cmd.extend(["--reasoning-budget-message", reasoning_budget_message])
+                self._reasoning_budget = reasoning_budget
+                self._reasoning_budget_message = reasoning_budget_message
+
                 if launch_mmproj_path and effective_is_vision:
                     cmd.extend(["--mmproj", launch_mmproj_path])
                     logger.info(f"Using mmproj for vision: {launch_mmproj_path}")
@@ -9098,6 +9128,10 @@ def _restore_after_tensor_downgrade():
 
                 # Library paths so llama-server finds its shared libs and CUDA DLLs.
                 env = self._llama_server_env_for_binary(binary)
+                # These launch settings are first-class and always have explicit
+                # defaults, so inherited llama.cpp env values must not contradict
+                # the command or the state echoed through /load and /status.
+                env.pop("LLAMA_ARG_THINK_BUDGET_MESSAGE", None)
                 if gpu_memory_mode == "manual":
                     self._clear_manual_placement_env(env)
                 # Omitting --threads relies on llama.cpp's physical-core default, so
@@ -10068,6 +10102,8 @@ def _already_in_target_state(
         is_vision: bool,
         gguf_path: Optional[str] = None,
         spec_draft_n_max: Optional[int] = None,
+        reasoning_budget: int = -1,
+        reasoning_budget_message: str = "",
         tensor_parallel: bool = False,
         gpu_memory_mode: Literal["auto", "manual"] = "auto",
         gpu_layers: int = -1,
@@ -10111,6 +10147,12 @@ def _norm(value):
 
         if _norm(self._cache_type_kv) != _norm(cache_type_kv):
             return False
+        if self._reasoning_budget != resolve_reasoning_budget(extra_args, reasoning_budget):
+            return False
+        if self._reasoning_budget_message != resolve_reasoning_budget_message(
+            extra_args, reasoning_budget_message
+        ):
+            return False
         # Reconcile a user --split-mode in extras AND an inherited tensor
         # LLAMA_ARG_SPLIT_MODE env, but only against a server that actually
         # launched tensor: if load_model downgraded to layer split it scrubbed
@@ -10356,6 +10398,8 @@ def unload_model(self) -> bool:
             self._chat_template_override = None
             self._supports_reasoning = False
             self._reasoning_always_on = False
+            self._reasoning_budget = -1
+            self._reasoning_budget_message = ""
             self._reasoning_style = "enable_thinking"
             self._reasoning_effort_levels = []
             self._reasoning_default = True

studio/backend/core/inference/llama_server_args.py

@@ -14,7 +14,9 @@
 from __future__ import annotations
 
 import os
-from typing import Iterable, Mapping, Optional
+from typing import Callable, Iterable, Mapping, Optional
+
+from utils.reasoning_budget import validate_reasoning_budget_message
 
 # Valid llama-server --parallel range, shared with LoadRequest.n_parallel.
 # Mirrored by callers that cannot import this: run.py and unsloth_cli/commands/
@@ -129,6 +131,8 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
     parse_cache_override(out)
     parse_split_mode_override(out)
     parse_gpu_layers_override(out)
+    parse_reasoning_budget_override(out)
+    parse_reasoning_budget_message_override(out)
     return out
 
 
@@ -146,6 +150,9 @@ def is_managed_flag(flag: str) -> bool:
 _CACHE_TYPE_K_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k"})
 _CACHE_TYPE_V_FLAGS: frozenset[str] = frozenset({"-ctv", "--cache-type-v"})
 _CACHE_FLAGS: frozenset[str] = _CACHE_TYPE_K_FLAGS | _CACHE_TYPE_V_FLAGS
+_REASONING_BUDGET_FLAGS: frozenset[str] = frozenset({"--reasoning-budget"})
+_REASONING_BUDGET_MESSAGE_FLAGS: frozenset[str] = frozenset({"--reasoning-budget-message"})
+_REASONING_BUDGET_MAX = 2_147_483_647
 _SPEC_FLAGS: frozenset[str] = frozenset(
     {
         "--spec-default",
@@ -278,7 +285,13 @@ def resolve_requested_ctx(args: Optional[Iterable[str]], fallback_n_ctx: int) ->
     return override if override is not None else fallback_n_ctx
 
 
-def _last_flag_value(args: Optional[Iterable[str]], flags: frozenset[str]) -> Optional[str]:
+def _last_flag_value(
+    args: Optional[Iterable[str]],
+    flags: frozenset[str],
+    *,
+    preserve_raw: bool = False,
+    validate_value: Optional[Callable[[str], object]] = None,
+) -> Optional[str]:
     """Return the last-wins string value among ``flags`` in extras, or None.
 
     Handles both ``--flag=value`` and ``--flag value`` forms and raises if a
@@ -307,10 +320,12 @@ def _last_flag_value(args: Optional[Iterable[str]], flags: frozenset[str]) -> Op
             raw_value = tokens[i + 1]
             i += 2
 
-        value = str(raw_value).strip()
-        if not value:
+        raw_value = str(raw_value)
+        if not raw_value.strip():
             raise ValueError(f"llama-server flag '{flag}' requires a non-empty value")
-        override = value
+        if validate_value is not None:
+            validate_value(raw_value)
+        override = raw_value if preserve_raw else raw_value.strip()
 
     return override
 
@@ -326,6 +341,49 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
     return _last_flag_value(args, _CACHE_FLAGS)
 
 
+def _validate_reasoning_budget_value(raw_value: str) -> int:
+    try:
+        value = int(raw_value)
+    except ValueError as exc:
+        raise ValueError("llama-server --reasoning-budget requires an integer value") from exc
+    if value < -1:
+        raise ValueError("llama-server --reasoning-budget requires a value of at least -1")
+    if value > _REASONING_BUDGET_MAX:
+        raise ValueError(
+            f"llama-server --reasoning-budget requires a value of at most {_REASONING_BUDGET_MAX}"
+        )
+    return value
+
+
+def parse_reasoning_budget_override(args: Optional[Iterable[str]]) -> Optional[int]:
+    """Return the last user-supplied ``--reasoning-budget`` value."""
+    raw_value = _last_flag_value(
+        args, _REASONING_BUDGET_FLAGS, validate_value = _validate_reasoning_budget_value
+    )
+    return None if raw_value is None else int(raw_value)
+
+
+def parse_reasoning_budget_message_override(args: Optional[Iterable[str]]) -> Optional[str]:
+    """Return the last user-supplied ``--reasoning-budget-message`` value."""
+    value = _last_flag_value(
+        args,
+        _REASONING_BUDGET_MESSAGE_FLAGS,
+        preserve_raw = True,
+        validate_value = validate_reasoning_budget_message,
+    )
+    return value
+
+
+def resolve_reasoning_budget(args: Optional[Iterable[str]], fallback: int) -> int:
+    override = parse_reasoning_budget_override(args)
+    return override if override is not None else fallback
+
+
+def resolve_reasoning_budget_message(args: Optional[Iterable[str]], fallback: str) -> str:
+    override = parse_reasoning_budget_message_override(args)
+    return override if override is not None else fallback
+
+
 def parse_gpu_layers_override(args: Optional[Iterable[str]]) -> Optional[int]:
     """Return the last user-supplied GPU layer count from extras.
 
@@ -482,6 +540,8 @@ def strip_shadowing_flags(
     strip_tensor_split: bool = False,
     strip_offload: bool = False,
     strip_device: bool = False,
+    strip_reasoning_budget: bool = False,
+    strip_reasoning_budget_message: bool = False,
 ) -> list[str]:
     """Strip flags that shadow first-class Unsloth settings.
 
@@ -515,6 +575,10 @@ def strip_shadowing_flags(
         shadowing |= _OFFLOAD_SHADOWING_FLAGS
     if strip_device:
         shadowing |= _DEVICE_FLAGS
+    if strip_reasoning_budget:
+        shadowing |= _REASONING_BUDGET_FLAGS
+    if strip_reasoning_budget_message:
+        shadowing |= _REASONING_BUDGET_MESSAGE_FLAGS
 
     tokens = [str(a) for a in (args or [])]
     out: list[str] = []

studio/backend/models/inference.py

@@ -20,6 +20,7 @@
 
 from core.inference.llama_server_args import PARALLEL_MAX, PARALLEL_MIN
 from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
+from utils.reasoning_budget import validate_reasoning_budget_message
 
 
 class LoadRequest(BaseModel):
@@ -205,6 +206,22 @@ def _reject_degenerate_tensor_split(cls, value: Optional[List[float]]) -> Option
             "auth, UI/server mode) are rejected. Ignored for non-GGUF models."
         ),
     )
+    reasoning_budget: int = Field(
+        -1,
+        ge = -1,
+        le = 2_147_483_647,
+        description = "llama-server reasoning token budget (-1 = model/default behavior).",
+    )
+    reasoning_budget_message: str = Field(
+        "",
+        description = "Message emitted by llama-server when the reasoning budget is exhausted.",
+    )
+
+    @field_validator("reasoning_budget_message")
+    @classmethod
+    def _validate_reasoning_budget_message(cls, value: str) -> str:
+        return validate_reasoning_budget_message(value)
+
     force_cancel_active: bool = Field(
         False,
         description = (
@@ -489,6 +506,10 @@ class LoadResponse(BaseModel):
         False,
         description = "Whether reasoning is always on (hardcoded <think> tags, not toggleable)",
     )
+    reasoning_budget: int = Field(-1, description = "Effective llama-server reasoning token budget.")
+    reasoning_budget_message: str = Field(
+        "", description = "Effective llama-server reasoning-budget exhaustion message."
+    )
     supports_preserve_thinking: bool = Field(
         False,
         description = "Whether the template understands the optional preserve_thinking kwarg (Qwen3.6-style)",
@@ -664,6 +685,10 @@ class InferenceStatusResponse(BaseModel):
     reasoning_always_on: bool = Field(
         False, description = "Whether reasoning is always on (not toggleable)"
     )
+    reasoning_budget: int = Field(-1, description = "Effective llama-server reasoning token budget.")
+    reasoning_budget_message: str = Field(
+        "", description = "Effective llama-server reasoning-budget exhaustion message."
+    )
     supports_preserve_thinking: bool = Field(
         False,
         description = "Whether the active model's template understands the optional preserve_thinking kwarg",

studio/backend/routes/chat_history.py

@@ -8,11 +8,12 @@
 from typing import Annotated, Any, Literal, Optional
 
 from fastapi import APIRouter, Depends, HTTPException, Query, Request
-from pydantic import BaseModel, ConfigDict, Field, ValidationError
+from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator
 
 from auth.authentication import get_current_subject
 from core.inference.llama_server_args import PARALLEL_MAX, PARALLEL_MIN
 from loggers import get_logger
+from utils.reasoning_budget import validate_reasoning_budget_message
 from utils.utils import safe_curated_detail, log_and_http_error
 from storage.studio_db import (
     ChatMessageConflictError,
@@ -171,11 +172,18 @@ class ChatPresetLoadConfig(BaseModel):
     speculativeType: Optional[str] = None
     specDraftNMax: Optional[int] = Field(default = None, ge = 1, le = 16)
     nParallel: Optional[int] = Field(default = None, ge = PARALLEL_MIN, le = PARALLEL_MAX)
+    reasoningBudget: Optional[int] = Field(default = None, ge = -1, le = 2_147_483_647)
+    reasoningBudgetMessage: Optional[str] = None
     tensorParallel: Optional[bool] = None
     gpuMemoryMode: Optional[Literal["manual"]] = None
     gpuLayers: Optional[int] = None
     nCpuMoe: Optional[int] = Field(default = None, ge = 0)
 
+    @field_validator("reasoningBudgetMessage")
+    @classmethod
+    def _validate_reasoning_budget_message(cls, value: Optional[str]) -> Optional[str]:
+        return None if value is None else validate_reasoning_budget_message(value)
+
 
 class ChatPreset(BaseModel):
     model_config = ConfigDict(extra = "forbid")