ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Pop the monitor open only for the traffic a caller made

Browse source 
Lifecycle rows are shared so a load or a download shows up in everyone's
monitor list, which is deliberate. Since they started carrying via_api_key
they also carry the flag the floating panel auto-opens on, and that reached
every authenticated subject: another logged-in browser sprang open for API
traffic it had nothing to do with. The row now records the caller that drove
it and reports the attribution only to them. Visibility is untouched, so the
row still appears for everybody, and a subject-scoped Clear hides a shared row
it owns rather than deleting it out of everyone else's history.

Auto-download had the flag hardcoded on, reasoning that only an API request
gets that far. Only a /v1 request does, which is not the same thing: Studio's
own chat calls those same endpoints with a session JWT, so a chat that named a
model this server does not have popped the panel open mid-chat, which is
exactly what via_api_key exists to prevent. The attribution now comes from the
request that asked for the download.
This commit is contained in:
danielhanchen 2026-07-29 00:18:42 +00:00
commit 69ea4e3136
5 changed files with 200 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

61
studio/backend/core/inference/api_monitor.py
View file

@ -49,6 +49,8 @@ class ApiMonitorEntry:
status: str
started_at: float
updated_at: float
# Who this row belongs to. On a shared lifecycle row it does not restrict
# visibility (see _visible); it names the caller the row is attributed to.
subject: Optional[str] = None
# True for sk-unsloth key callers, not UI sessions. The floating panel only
# auto-opens for these, so Studio's own chat does not pop it mid-chat.
@ -72,7 +74,12 @@ class ApiMonitorEntry:
# 0-100 for a running download row; None when not applicable.
progress: Optional[float] = None
def snapshot(self, *, include_details: bool = True) -> dict[str, Any]:
def snapshot(
self,
*,
include_details: bool = True,
attributed: bool = True,
) -> dict[str, Any]:
duration_ms = None
if self.finished_monotonic is not None:
duration_ms = max(
@ -89,7 +96,10 @@ class ApiMonitorEntry:
"endpoint": self.endpoint,
"method": self.method,
"model": self.model,
"via_api_key": self.via_api_key,
# A lifecycle row is shared, so it reaches subjects that had nothing to
# do with it. The overlay auto-opens on this flag, so report it only to
# the caller the row is attributed to; everyone else still sees the row.
"via_api_key": self.via_api_key and attributed,
"prompt_preview": _trim(self.prompt, _PREVIEW_CHARS),
"reply_preview": _trim(self.reply, _PREVIEW_CHARS),
"prompt_truncated": len(self.prompt) > _PREVIEW_CHARS,
@ -174,12 +184,17 @@ class ApiMonitor:
reason: Optional[str] = None,
running: bool = False,
via_api_key: bool = False,
subject: Optional[str] = None,
) -> str:
"""Record a model load/unload alongside the request traffic that caused it.
``running=True`` opens the row for the caller to close with :meth:`finish` /
:meth:`fail`; an unload is terminal on arrival. Rows are shared (visible to
every subject) and share the request retention budget.
``subject`` names the caller whose request drove this, which is what
``via_api_key`` is reported to; it does not narrow who sees the row. Pass it
whenever ``via_api_key`` is set, or the attribution reaches nobody.
"""
if not self._enabled:
return ""
@ -205,6 +220,9 @@ class ApiMonitor:
# trace of it, and without the attribution the monitor stayed shut on
# exactly the failures it exists to surface.
via_api_key = via_api_key,
# Shared rows are read by every subject, so the attribution needs an
# owner or the pop-open lands in browsers that did not cause it.
subject = subject,
)
with self._lock:
self._entries.appendleft(entry)
@ -371,7 +389,10 @@ class ApiMonitor:
) -> list[dict[str, Any]]:
with self._lock:
return [
entry.snapshot(include_details = include_details)
entry.snapshot(
include_details = include_details,
attributed = self._attributed(entry, subject),
)
for entry in self._entries
if self._visible(entry, subject)
]
@ -388,7 +409,10 @@ class ApiMonitor:
return None
if not self._visible(entry, subject):
return None
return entry.snapshot(include_details = True)
return entry.snapshot(
include_details = True,
attributed = self._attributed(entry, subject),
)
def active_count(self, *, subject: Optional[str] = None) -> int:
# Lifecycle rows show as "running" while loading but are not in-flight API requests.
@ -416,18 +440,35 @@ class ApiMonitor:
# A running shared row is a load in progress, not history, so it stays.
hidden = self._hidden_shared.setdefault(subject, set())
for entry in self._entries:
if entry.shared and entry.subject != subject and entry.status != "running":
if entry.shared and entry.status != "running":
hidden.add(entry.id)
self._entries = deque(entry for entry in self._entries if entry.subject != subject)
# Shared rows are hidden, never dropped, even when this subject owns
# one: they are another caller's history too, and an owned shared row
# is exactly what an API-key load produces.
self._entries = deque(
entry
for entry in self._entries
if entry.shared or entry.subject != subject
)
def _visible(self, entry: ApiMonitorEntry, subject: Optional[str]) -> bool:
if subject is None:
return True
if entry.subject == subject:
if entry.shared:
# Shared rows reach every subject, minus the ones this one cleared.
# Checked before ownership so clearing hides a subject's own rows too.
return entry.id not in self._hidden_shared.get(subject, ())
return entry.subject == subject
def _attributed(self, entry: ApiMonitorEntry, subject: Optional[str]) -> bool:
"""Whether *subject* is the caller this row's API traffic belongs to.
Only they should have the overlay pop open for it. An unscoped read (no
subject: internal callers and tests) sees the row's own flag.
"""
if subject is None:
return True
if not entry.shared:
return False
return entry.id not in self._hidden_shared.get(subject, ())
return entry.subject == subject
def _find_locked(self, entry_id: str) -> Optional[ApiMonitorEntry]:
for entry in self._entries:

42
studio/backend/core/inference/openai_auto_download.py
View file

@ -448,6 +448,8 @@ async def maybe_auto_download(
*,
hf_token: Optional[str] = None,
require_vision: bool = False,
subject: Optional[str] = None,
via_api_key: bool = False,
) -> Optional[AutoDownloadRefusal]:
"""Start (or report on) a background fetch of *requested_model*.
@ -457,6 +459,10 @@ async def maybe_auto_download(
``require_vision`` refuses a target with no mmproj companion rather than spend
gigabytes on weights that cannot answer the request; the local capability guard
only ever sees an already-downloaded model.
``subject`` and ``via_api_key`` describe the caller for the monitor row this
opens: the same /v1 endpoints serve Studio's own chat on a session JWT, so the
download is not API-key traffic unless the request that asked for it was.
"""
global _active
@ -529,7 +535,14 @@ async def maybe_auto_download(
try:
return await _admit_and_start(
repo_id, wanted_variant, requested_model, hf_token, provisional, require_vision
repo_id,
wanted_variant,
requested_model,
hf_token,
provisional,
require_vision,
subject = subject,
via_api_key = via_api_key,
)
except BaseException:
# Not `except Exception`: a cancel mid-probe would otherwise wedge the provisional slot.
@ -544,6 +557,9 @@ async def _admit_and_start(
hf_token: Optional[str],
active: _Active,
require_vision: bool = False,
*,
subject: Optional[str] = None,
via_api_key: bool = False,
) -> Optional[AutoDownloadRefusal]:
from hub.utils.hf_errors import hf_error_status
@ -690,7 +706,16 @@ async def _admit_and_start(
),
)
return await _dispatch(repo_id, variant, expected_bytes, requested_model, hf_token, active)
return await _dispatch(
repo_id,
variant,
expected_bytes,
requested_model,
hf_token,
active,
subject = subject,
via_api_key = via_api_key,
)
def preferred_quant(labels) -> Optional[str]:
@ -737,6 +762,9 @@ async def _dispatch(
requested_model: str,
hf_token: Optional[str],
active: _Active,
*,
subject: Optional[str] = None,
via_api_key: bool = False,
) -> AutoDownloadRefusal:
global _active
@ -777,12 +805,18 @@ async def _dispatch(
return busy
monitor_id = api_monitor.record_lifecycle(
# Only an API request reaches auto-download, hence reason "api".
# Reason "api" because only a /v1 request reaches auto-download. That is not
# the same as API-key traffic though: Studio's own chat calls those same
# endpoints with a session JWT, and marking its download as API traffic pops
# the overlay mid-chat, which via_api_key exists to prevent. So take the
# attribution from the request, and name the caller it belongs to, since the
# row is shared with every other subject.
event = "download",
model = label,
reason = "api",
running = True,
via_api_key = True,
via_api_key = via_api_key,
subject = subject,
)
with _lock:
if _active is active:

13
studio/backend/routes/inference.py
View file

@ -3855,6 +3855,7 @@ async def _maybe_auto_download_model(
fastapi_request: Optional[Request],
*,
require_vision: bool = False,
current_subject: Optional[str] = None,
) -> None:
"""Opt-in: start fetching a named GGUF this server doesn't have.
@ -3877,6 +3878,10 @@ async def _maybe_auto_download_model(
requested_model,
hf_token = _auto_download_hf_token(fastapi_request),
require_vision = require_vision,
subject = current_subject,
# These /v1 endpoints also serve Studio's own chat on a session JWT,
# so only mark the download row as API traffic when it really is.
via_api_key = _request_used_api_key(fastapi_request),
)
except Exception as exc:
# Never turn a servable request into a 500 over the download attempt.
@ -4246,7 +4251,10 @@ async def _maybe_auto_switch_model(
# Not on disk. Opt-in: fetch in the background and ask the caller to retry.
if auto_switch_on and not reload_only:
await _maybe_auto_download_model(
requested_model, fastapi_request, require_vision = require_vision
requested_model,
fastapi_request,
require_vision = require_vision,
current_subject = current_subject,
)
# Idle-unload may have freed the model; reload exactly what it freed
# (path + quant + advertised id) so an alias/unknown name stays servable
@ -5324,6 +5332,9 @@ async def _load_model_impl(
# Auto-switch loads run before the endpoint opens its request row, so a
# load that fails there leaves this as the only trace of API traffic.
via_api_key = _request_used_api_key(fastapi_request),
# The row is shared, so every subject reads it. Name the caller it belongs
# to or the overlay pops open in browsers that did not cause the traffic.
subject = current_subject,
)
native_grant_backed = False

45
studio/backend/tests/test_api_monitor.py
View file

@ -542,3 +542,48 @@ def test_an_api_triggered_lifecycle_row_carries_the_attribution():
after = {e["id"]: e for e in monitor.snapshot()}
assert after[api_load]["via_api_key"] is True
assert after[api_load]["status"] == "error"
def test_an_api_lifecycle_row_pops_the_overlay_only_for_its_own_caller():
"""A lifecycle row is shared so it appears in every monitor list, and it also
carries via_api_key, which is what the floating panel auto-opens on. Reported
to everyone, the panel springs open in a browser that had nothing to do with
the traffic. The row stays visible to all; only the attribution is scoped."""
monitor = ApiMonitor(max_entries = 5)
row = monitor.record_lifecycle(
event = "load",
model = "org/Repo-GGUF",
running = True,
via_api_key = True,
subject = "alice",
)
mine = {e["id"]: e for e in monitor.snapshot(subject = "alice")}
theirs = {e["id"]: e for e in monitor.snapshot(subject = "bob")}
# Shared visibility is deliberate and must survive: bob still sees the load.
assert row in mine and row in theirs
assert mine[row]["via_api_key"] is True
assert theirs[row]["via_api_key"] is False
# The details read is scoped the same way, so the panel cannot re-derive it.
assert monitor.get(row, subject = "alice")["via_api_key"] is True
assert monitor.get(row, subject = "bob")["via_api_key"] is False
# An unscoped read (internal callers) still sees the row's own flag.
assert monitor.get(row)["via_api_key"] is True
def test_clearing_hides_a_shared_row_this_caller_owns_rather_than_deleting_it():
"""An API-key load now owns its shared row. A subject-scoped clear drops that
subject's rows, so without this the owner's Clear would delete a row every
other caller can still see and wipe it out of their history too."""
monitor = ApiMonitor(max_entries = 10)
row = monitor.record_lifecycle(
event = "unload", model = "org/Repo-GGUF", via_api_key = True, subject = "alice"
)
monitor.clear(subject = "alice")
assert monitor.snapshot(subject = "alice") == []
assert {e["id"] for e in monitor.snapshot(subject = "bob")} == {row}
assert monitor.get(row, subject = "bob") is not None

57
studio/backend/tests/test_openai_auto_download.py
View file

@ -684,13 +684,17 @@ class _Req:
self.headers = headers or {}
def _hook(model, request, enabled):
def _hook(model, request, enabled, current_subject = None):
import utils.openai_auto_switch_settings as s
original = s.get_openai_auto_download_enabled
s.get_openai_auto_download_enabled = lambda: enabled
try:
return asyncio.run(inference_route._maybe_auto_download_model(model, request))
return asyncio.run(
inference_route._maybe_auto_download_model(
model, request, current_subject = current_subject
)
)
finally:
s.get_openai_auto_download_enabled = original
@ -726,13 +730,60 @@ def test_hook_uses_the_anthropic_envelope_on_messages(hub):
def test_hook_swallows_unexpected_failures(hub, monkeypatch):
# A broken download path must not turn a servable request into a 500.
async def _boom(model, hf_token = None):
async def _boom(model, hf_token = None, **kwargs):
raise RuntimeError("boom")
monkeypatch.setattr(auto_dl, "maybe_auto_download", _boom)
assert _hook("unsloth/x-GGUF", _Req(), enabled = True) is None
def _download_rows():
from core.inference.api_monitor import api_monitor
return [e for e in api_monitor.snapshot() if e["event"] == "download"]
def test_a_ui_session_download_is_not_marked_as_api_traffic(hub):
"""The monitor overlay auto-opens on via_api_key, which exists to separate
"someone is serving other clients" from "someone is using Unsloth". Studio's
own chat hits these same /v1 endpoints with a session JWT, so hardcoding the
flag on the download row popped the panel open mid-chat."""
from fastapi import HTTPException
from core.inference.api_monitor import api_monitor
api_monitor.clear()
with pytest.raises(HTTPException):
# No Authorization header: the UI's session-JWT path.
_hook("unsloth/x-GGUF", _Req(), enabled = True, current_subject = "unsloth")
rows = _download_rows()
assert rows and all(row["via_api_key"] is False for row in rows)
def test_an_api_key_download_keeps_the_attribution_and_names_its_caller(hub):
"""The row is shared, so it needs the subject as well: without one the
attribution is reported to every logged-in browser instead of the caller."""
from fastapi import HTTPException
from auth.authentication import API_KEY_PREFIX
from core.inference.api_monitor import api_monitor
api_monitor.clear()
with pytest.raises(HTTPException):
_hook(
"unsloth/x-GGUF",
_Req(headers = {"authorization": f"Bearer {API_KEY_PREFIX}abc123"}),
enabled = True,
current_subject = "unsloth",
)
rows = _download_rows()
assert rows and all(row["via_api_key"] is True for row in rows)
# Still shared: another subject sees the row, just not the attribution.
others = [
e for e in api_monitor.snapshot(subject = "someone-else") if e["event"] == "download"
]
assert len(others) == len(rows)
assert all(row["via_api_key"] is False for row in others)
def test_hook_prefers_the_hub_header_token(hub):
from fastapi import HTTPException
from hub.dependencies import HUB_HF_TOKEN_HEADER