diff --git a/studio/backend/core/inference/api_monitor.py b/studio/backend/core/inference/api_monitor.py index 10d177fc23..92b4f43879 100644 --- a/studio/backend/core/inference/api_monitor.py +++ b/studio/backend/core/inference/api_monitor.py @@ -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: diff --git a/studio/backend/core/inference/openai_auto_download.py b/studio/backend/core/inference/openai_auto_download.py index 2435841312..7252d3382d 100644 --- a/studio/backend/core/inference/openai_auto_download.py +++ b/studio/backend/core/inference/openai_auto_download.py @@ -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: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 69362486dd..7a3f02250f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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 diff --git a/studio/backend/tests/test_api_monitor.py b/studio/backend/tests/test_api_monitor.py index 13e6b8f39e..2ad5eb2397 100644 --- a/studio/backend/tests/test_api_monitor.py +++ b/studio/backend/tests/test_api_monitor.py @@ -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 diff --git a/studio/backend/tests/test_openai_auto_download.py b/studio/backend/tests/test_openai_auto_download.py index b1dde175da..f14ab02ad9 100644 --- a/studio/backend/tests/test_openai_auto_download.py +++ b/studio/backend/tests/test_openai_auto_download.py @@ -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