Merge branch 'main' into feature/rag

studio/backend/core/inference/llama_cpp.py

@@ -28,6 +28,12 @@ from urllib.parse import urlparse
 
 import httpx
 
+from core.inference.llama_server_args import (
+    parse_cache_override,
+    parse_ctx_override,
+    resolve_cache_type_kv,
+    resolve_requested_ctx,
+)
 from core.tool_healing import (
     _TC_END_TAG_RE,
     _TC_FUNC_CLOSE_RE,
@@ -2737,7 +2743,23 @@ class LlamaCppBackend:
                 # Select GPU(s) based on model size + estimated KV cache.
                 # Seed safe defaults before GPU probing so the except path
                 # still has valid state to publish.
-                effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0)
+                ctx_override = parse_ctx_override(extra_args)
+                requested_ctx = resolve_requested_ctx(extra_args, n_ctx)
+                cache_override = parse_cache_override(extra_args)
+                cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv)
+                if ctx_override is not None and ctx_override > 0:
+                    logger.info(
+                        f"User --ctx-size {ctx_override} honored; "
+                        "skipping auto-reduce"
+                    )
+                if cache_override is not None:
+                    logger.info(
+                        f"User --cache-type-k/-v {cache_override} "
+                        "honored for KV estimate"
+                    )
+                effective_ctx = (
+                    requested_ctx if requested_ctx > 0 else (self._context_length or 0)
+                )
                 max_available_ctx = self._context_length or effective_ctx
                 gpus: list[tuple[int, int]] = []
                 try:
@@ -2747,8 +2769,8 @@ class LlamaCppBackend:
                     # Resolve effective context: 0 means let llama-server use the
                     # model's native length.  Only expand to a known native length
                     # if metadata is available; otherwise preserve 0 as a sentinel.
-                    if n_ctx > 0:
-                        effective_ctx = n_ctx
+                    if requested_ctx > 0:
+                        effective_ctx = requested_ctx
                     elif self._context_length is not None:
                         effective_ctx = self._context_length
                     else:
@@ -2801,7 +2823,7 @@ class LlamaCppBackend:
                     #   since multi-GPU is slower and the user didn't ask for a
                     #   specific context length.
                     gpu_indices, use_fit = None, True
-                    explicit_ctx = n_ctx > 0
+                    explicit_ctx = requested_ctx > 0
 
                     if gpus and self._can_estimate_kv() and effective_ctx > 0:
                         # Compute the largest hardware-aware cap from the model's
@@ -2858,7 +2880,7 @@ class LlamaCppBackend:
                             gpu_indices, use_fit = self._select_gpus(
                                 requested_total, gpus
                             )
-                            # No silent shrink: effective_ctx stays == n_ctx.
+                            # No silent shrink: effective_ctx stays == requested_ctx.
                         else:
                             # Auto context: prefer fewer GPUs, cap context
                             # to fit. Same headroom threshold as
@@ -2947,7 +2969,7 @@ class LlamaCppBackend:
                 except Exception as e:
                     logger.warning(f"GPU selection failed ({e}), using --fit on")
                     gpu_indices, use_fit = None, True
-                    effective_ctx = n_ctx  # fall back to original
+                    effective_ctx = requested_ctx  # fall back to original
 
                 launch_mmproj_path = self._resolve_launch_mmproj_path(
                     model_path = model_path,

studio/backend/core/inference/llama_server_args.py

@@ -110,6 +110,8 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
                 f"and cannot be passed as an extra arg"
             )
         out.append(token)
+    parse_ctx_override(out)
+    parse_cache_override(out)
     return out
 
 
@@ -166,6 +168,117 @@ _BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset(
 )
 
 
+def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]:
+    """Return the last user-supplied ``-c`` / ``--ctx-size`` value.
+
+    Mirrors llama.cpp's last-wins flag parsing for the one pass-through
+    numeric knob Studio's load-time fit logic needs to see.
+    """
+    if not args:
+        return None
+
+    tokens = [str(a) for a in args]
+    override: Optional[int] = None
+    i, n = 0, len(tokens)
+    while i < n:
+        tok = tokens[i]
+        flag = _flag_name(tok)
+        if flag is None or flag not in _CONTEXT_FLAGS:
+            i += 1
+            continue
+
+        if "=" in tok:
+            raw_value = tok.split("=", 1)[1]
+            i += 1
+        else:
+            if i + 1 >= n or _flag_name(tokens[i + 1]) is not None:
+                raise ValueError(
+                    f"llama-server flag '{flag}' requires an integer value"
+                )
+            raw_value = tokens[i + 1]
+            i += 2
+
+        try:
+            value = int(str(raw_value).strip())
+        except ValueError as exc:
+            raise ValueError(
+                f"llama-server flag '{flag}' requires an integer value"
+            ) from exc
+        if value < 0:
+            raise ValueError(
+                f"llama-server flag '{flag}' requires a non-negative integer value"
+            )
+        override = value
+
+    return override
+
+
+def resolve_requested_ctx(
+    args: Optional[Iterable[str]],
+    fallback_n_ctx: int,
+) -> int:
+    """Return the context size load_model should treat as requested.
+
+    Single source of truth for the two-line ``ctx_override = parse_ctx_override(...);
+    requested_ctx = ctx_override if ctx_override is not None else n_ctx`` pattern
+    used by ``load_model`` so tests don't have to reimplement the conditional
+    locally and then assert against their own reimplementation.
+    """
+    override = parse_ctx_override(args)
+    return override if override is not None else fallback_n_ctx
+
+
+def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
+    """Return the last-wins cache type if extras pass cache flags.
+
+    Mirrors parse_ctx_override but for cache type. Recognises both -ctk
+    (key) and -ctv (value). When both flags appear, returns the last-wins
+    value, treating key and value cache flags as the same setting because
+    Studio's KV estimate has a single cache_type_kv knob.
+    """
+    if not args:
+        return None
+
+    tokens = [str(a) for a in args]
+    override: Optional[str] = None
+    i, n = 0, len(tokens)
+    while i < n:
+        tok = tokens[i]
+        flag = _flag_name(tok)
+        if flag is None or flag not in _CACHE_FLAGS:
+            i += 1
+            continue
+
+        if "=" in tok:
+            raw_value = tok.split("=", 1)[1]
+            i += 1
+        else:
+            if i + 1 >= n or _flag_name(tokens[i + 1]) is not None:
+                raise ValueError(f"llama-server flag '{flag}' requires a value")
+            raw_value = tokens[i + 1]
+            i += 2
+
+        value = str(raw_value).strip()
+        if not value:
+            raise ValueError(f"llama-server flag '{flag}' requires a non-empty value")
+        override = value
+
+    return override
+
+
+def resolve_cache_type_kv(
+    args: Optional[Iterable[str]],
+    fallback_cache_type_kv: Optional[str],
+) -> Optional[str]:
+    """Return the cache type load_model should treat as requested.
+
+    Single source of truth for the cache override conditional used by
+    ``load_model``.
+    """
+    override = parse_cache_override(args)
+    return override if override is not None else fallback_cache_type_kv
+
+
 def strip_shadowing_flags(
     args: Iterable[str],
     *,

studio/backend/routes/inference.py

@@ -3137,9 +3137,12 @@ async def openai_chat_completions(
         else:
             try:
                 full_text = ""
+                completion_usage = None
                 for token in gguf_generate():
                     if isinstance(token, dict):
-                        continue  # skip metadata dict in non-streaming path
+                        if token.get("type") == "metadata":
+                            completion_usage = token.get("usage")
+                        continue
                     full_text = token
 
                 response = ChatCompletion(
@@ -3152,6 +3155,15 @@ async def openai_chat_completions(
                             finish_reason = "stop",
                         )
                     ],
+                    usage = CompletionUsage(
+                        prompt_tokens = (completion_usage or {}).get("prompt_tokens")
+                        or 0,
+                        completion_tokens = (completion_usage or {}).get(
+                            "completion_tokens"
+                        )
+                        or 0,
+                        total_tokens = (completion_usage or {}).get("total_tokens") or 0,
+                    ),
                 )
                 return JSONResponse(content = response.model_dump())
 

studio/backend/tests/test_gguf_completion_usage.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+"""Regression tests for GGUF non-streaming chat completion usage."""
+
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from auth.authentication import get_current_subject
+import routes.inference as inference_route
+
+
+class _GgufBackend:
+    is_loaded = True
+    model_identifier = "test/model.gguf"
+    _is_audio = False
+    is_vision = False
+    supports_tools = False
+
+    def __init__(self, usage):
+        self.usage = usage
+
+    def generate_chat_completion(self, **kwargs):
+        yield "answer"
+        yield {
+            "type": "metadata",
+            "usage": self.usage,
+            "timings": {"prompt_n": 23, "predicted_n": 1283},
+        }
+
+
+def _request_completion(monkeypatch, usage):
+    monkeypatch.setattr(
+        inference_route, "get_llama_cpp_backend", lambda: _GgufBackend(usage)
+    )
+    monkeypatch.setattr(
+        inference_route, "_effective_enable_tools", lambda payload: False
+    )
+
+    app = FastAPI()
+    app.include_router(inference_route.router)
+    app.dependency_overrides[get_current_subject] = lambda: "test-user"
+
+    return TestClient(app).post(
+        "/chat/completions",
+        json = {
+            "messages": [{"role": "user", "content": "Why is the sky blue?"}],
+            "stream": False,
+        },
+    )
+
+
+def test_non_streaming_gguf_completion_includes_generated_usage(monkeypatch):
+    response = _request_completion(
+        monkeypatch,
+        {"prompt_tokens": 23, "completion_tokens": 1283, "total_tokens": 1306},
+    )
+
+    assert response.status_code == 200
+    assert response.json()["usage"] == {
+        "prompt_tokens": 23,
+        "completion_tokens": 1283,
+        "total_tokens": 1306,
+    }
+
+
+def test_non_streaming_gguf_completion_defaults_nullable_usage_to_zero(monkeypatch):
+    response = _request_completion(
+        monkeypatch,
+        {"prompt_tokens": None, "completion_tokens": 1283, "total_tokens": None},
+    )
+
+    assert response.status_code == 200
+    assert response.json()["usage"] == {
+        "prompt_tokens": 0,
+        "completion_tokens": 1283,
+        "total_tokens": 0,
+    }

studio/backend/tests/test_llama_cpp_context_fit.py

@@ -84,6 +84,7 @@ _httpx_stub.Client = type(
 sys.modules.setdefault("httpx", _httpx_stub)
 
 from core.inference.llama_cpp import LlamaCppBackend
+from core.inference.llama_server_args import parse_ctx_override, resolve_requested_ctx
 
 
 # ---------------------------------------------------------------------------
@@ -131,6 +132,7 @@ def _drive(
     native_ctx = 131072,
     kv_per_token_bytes = 325_000,
     can_estimate_kv = True,
+    extra_args = None,
 ):
     """Drive the post-metadata portion of load_model with stubbed inputs.
 
@@ -148,11 +150,16 @@ def _drive(
     inst._can_estimate_kv = lambda: can_estimate_kv
 
     context_length = inst._context_length
+    # Use the production helper instead of reimplementing the conditional
+    # locally; reimplementing makes the test pass for the test's own logic
+    # rather than production's, and silent drift won't be caught.
+    ctx_override = parse_ctx_override(extra_args)
+    requested_ctx = resolve_requested_ctx(extra_args, n_ctx)
 
-    effective_ctx = n_ctx if n_ctx > 0 else (context_length or 0)
+    effective_ctx = requested_ctx if requested_ctx > 0 else (context_length or 0)
     max_available_ctx = context_length or effective_ctx
-    if n_ctx > 0:
-        effective_ctx = n_ctx
+    if requested_ctx > 0:
+        effective_ctx = requested_ctx
     elif context_length is not None:
         effective_ctx = context_length
     else:
@@ -161,7 +168,7 @@ def _drive(
     max_available_ctx = context_length or effective_ctx
 
     gpu_indices, use_fit = None, True
-    explicit_ctx = n_ctx > 0
+    explicit_ctx = requested_ctx > 0
 
     if gpus and inst._can_estimate_kv() and effective_ctx > 0:
         native_ctx_for_cap = context_length or effective_ctx
@@ -236,6 +243,7 @@ def _drive(
         "gpu_indices": gpu_indices,
         "max_available_ctx": max_available_ctx,
         "original_ctx": original_ctx,
+        "ctx_override": ctx_override,
     }
 
 
@@ -349,6 +357,48 @@ class TestExplicitCtxRespectsUser:
         assert plan["c_arg"] == 2048
 
 
+# ---------------------------------------------------------------------------
+# Pass-through --ctx-size participates in context fit (#5676).
+# ---------------------------------------------------------------------------
+
+
+class TestExtraArgsCtxOverride:
+    def test_ctx_size_extra_honored_over_auto(self):
+        plan = _drive(
+            n_ctx = 0,
+            model_gib = 131,
+            gpus = [(0, 97_000)],
+            native_ctx = 196608,
+            extra_args = ["--ctx-size", "128000"],
+        )
+        assert plan["ctx_override"] == 128000
+        assert plan["original_ctx"] == 128000
+        assert plan["c_arg"] == 128000
+        assert plan["use_fit"] is True
+
+    def test_ctx_size_short_alias_honored_over_auto(self):
+        plan = _drive(
+            n_ctx = 0,
+            model_gib = 131,
+            gpus = [(0, 97_000)],
+            native_ctx = 196608,
+            extra_args = ["-c", "128000"],
+        )
+        assert plan["c_arg"] == 128000
+        assert plan["use_fit"] is True
+
+    def test_ctx_size_extra_wins_over_first_class_field(self):
+        plan = _drive(
+            n_ctx = 4096,
+            model_gib = 8,
+            gpus = [(0, 24_000)],
+            native_ctx = 131072,
+            extra_args = ["--ctx-size", "128000"],
+        )
+        assert plan["original_ctx"] == 128000
+        assert plan["c_arg"] == 128000
+
+
 # ---------------------------------------------------------------------------
 # Non-regression: fittable + auto still auto-picks largest fitting ctx
 # ---------------------------------------------------------------------------

studio/backend/tests/test_llama_server_args.py

@@ -30,6 +30,9 @@ _spec = importlib.util.spec_from_file_location("_lsa_test_only", _LSA_PATH)
 _lsa = importlib.util.module_from_spec(_spec)
 _spec.loader.exec_module(_lsa)
 is_managed_flag = _lsa.is_managed_flag
+parse_cache_override = _lsa.parse_cache_override
+parse_ctx_override = _lsa.parse_ctx_override
+resolve_cache_type_kv = _lsa.resolve_cache_type_kv
 strip_shadowing_flags = _lsa.strip_shadowing_flags
 validate_extra_args = _lsa.validate_extra_args
 
@@ -410,6 +413,86 @@ def test_is_managed_flag_false_for_mtp_pass_through():
     assert is_managed_flag("--spec-ngram-mod-n-max") is False
 
 
+# ── parse_ctx_override ───────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "args,expected",
+    [
+        (None, None),
+        ([], None),
+        (["--top-k", "20"], None),
+        (["--ctx-size", "128000"], 128000),
+        (["--ctx-size=128000"], 128000),
+        (["-c", "128000"], 128000),
+        (["-c=128000"], 128000),
+        (["-c", "4096", "--ctx-size", "128000"], 128000),
+    ],
+)
+def test_parse_ctx_override(args, expected):
+    assert parse_ctx_override(args) == expected
+
+
+@pytest.mark.parametrize(
+    "args",
+    [
+        ["--ctx-size"],
+        ["--ctx-size", "--top-k"],
+        ["--ctx-size", "abc"],
+        ["--ctx-size=abc"],
+        ["-c", "-1"],
+    ],
+)
+def test_parse_ctx_override_rejects_malformed_values(args):
+    with pytest.raises(ValueError, match = "ctx-size|'-c'"):
+        parse_ctx_override(args)
+
+
+def test_validate_extra_args_rejects_malformed_ctx_override():
+    with pytest.raises(ValueError, match = "ctx-size"):
+        validate_extra_args(["--ctx-size", "abc"])
+
+
+# ── parse_cache_override ─────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "args,expected",
+    [
+        (None, None),
+        ([], None),
+        (["--top-k", "20"], None),
+        (["--cache-type-k", "q8_0"], "q8_0"),
+        (["-ctk", "q4_0"], "q4_0"),
+        (["-ctv", "q4_0"], "q4_0"),
+        (["--cache-type-k=q4_0"], "q4_0"),
+        (["-ctk", "f16", "-ctk", "q8_0"], "q8_0"),
+    ],
+)
+def test_parse_cache_override(args, expected):
+    assert parse_cache_override(args) == expected
+
+
+@pytest.mark.parametrize(
+    "args",
+    [
+        ["-ctk"],
+        ["-ctk", "-c", "4096"],
+    ],
+)
+def test_parse_cache_override_rejects_malformed_values(args):
+    with pytest.raises(ValueError, match = "cache-type|'-ctk'"):
+        parse_cache_override(args)
+
+
+def test_resolve_cache_type_kv_uses_override_when_present():
+    assert resolve_cache_type_kv(["--cache-type-k", "q8_0"], "f16") == "q8_0"
+
+
+def test_resolve_cache_type_kv_uses_fallback_without_override():
+    assert resolve_cache_type_kv(["--top-k", "20"], "f16") == "f16"
+
+
 def test_strip_shadowing_flags_boolean_does_not_consume_next_token():
     # `--spec-default` is boolean; drop just the flag, keep the next token.
     out = strip_shadowing_flags(["--spec-default", "ngram-mod"], strip_spec = True)