mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-10-19 07:34:58 +00:00
fix(egress): gate the catalog's own model download, not just the manual pull
The Download button in the model catalog hits POST /llm/install, which
never called egress.require() — only the unused /providers/{provider}/pull
route did. Chat models installed through the catalog UI left silently,
with no consent popup and no destination logged. Also derive the actual
host from the resolved pull target, since Ollama's hf.co/<repo> fallback
installs go straight to huggingface.co, not registry.ollama.ai.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
parent
435f409df4
commit
48f68ad62c
6 changed files with 117 additions and 5 deletions
|
|
@ -18,9 +18,9 @@ HOSTS = {OLLAMA_PULL: "registry.ollama.ai", IMAGE_MODEL_PULL: "huggingface.co"}
|
|||
|
||||
|
||||
class EgressDeniedError(Exception):
|
||||
def __init__(self, destination: str) -> None:
|
||||
def __init__(self, destination: str, host: str | None = None) -> None:
|
||||
self.destination = destination
|
||||
self.host = host_of(destination)
|
||||
self.host = host or host_of(destination)
|
||||
super().__init__(f"sending data to {self.host} is off in Settings > Network")
|
||||
|
||||
|
||||
|
|
@ -28,6 +28,13 @@ def host_of(destination: str) -> str:
|
|||
return HOSTS.get(destination) or destination.removeprefix(HOST_PREFIX)
|
||||
|
||||
|
||||
def ollama_pull_host(model_name: str) -> str:
|
||||
"""Ollama fetches an `hf.co/<repo>` pull straight from Hugging Face, not its own registry."""
|
||||
if model_name.startswith("hf.co/"):
|
||||
return "huggingface.co"
|
||||
return HOSTS[OLLAMA_PULL]
|
||||
|
||||
|
||||
def is_destination(value: str) -> bool:
|
||||
return value in HOSTS or (
|
||||
value.startswith(HOST_PREFIX) and len(value) > len(HOST_PREFIX)
|
||||
|
|
@ -49,12 +56,14 @@ def _is_loopback(host: str) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def require(session: Session, destination: str | None) -> None:
|
||||
def require(
|
||||
session: Session, destination: str | None, host: str | None = None
|
||||
) -> None:
|
||||
if destination is None:
|
||||
return
|
||||
row = session.get(EgressDestination, destination)
|
||||
if row is None or not row.enabled:
|
||||
raise EgressDeniedError(destination)
|
||||
raise EgressDeniedError(destination, host)
|
||||
row.last_call_at = datetime.now(UTC)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from fastapi.responses import StreamingResponse
|
|||
from sqlalchemy.orm import Session
|
||||
|
||||
from api.dependencies import SessionDep, transact
|
||||
from modules.egress import service as egress
|
||||
from modules.llm.models import ModelRole, SelectedModel
|
||||
from modules.llm.recommendations.catalog import (
|
||||
InsufficientDiskError,
|
||||
|
|
@ -61,6 +62,7 @@ async def install_model(
|
|||
payload: InstallRequest,
|
||||
request: Request,
|
||||
service: CatalogServiceDep,
|
||||
session: SessionDep,
|
||||
) -> StreamingResponse:
|
||||
try:
|
||||
runtime, _model, plan = await service.preflight(payload.catalog_id)
|
||||
|
|
@ -81,6 +83,13 @@ async def install_model(
|
|||
except RuntimeError as error:
|
||||
raise HTTPException(status.HTTP_409_CONFLICT, str(error)) from error
|
||||
|
||||
await transact(
|
||||
session,
|
||||
egress.require,
|
||||
egress.OLLAMA_PULL,
|
||||
egress.ollama_pull_host(plan.model_name),
|
||||
)
|
||||
|
||||
lock = service.install_lock(runtime.name)
|
||||
if lock.locked():
|
||||
raise HTTPException(
|
||||
|
|
|
|||
|
|
@ -328,7 +328,12 @@ async def pull_model(
|
|||
service: CatalogServiceDep,
|
||||
session: SessionDep,
|
||||
) -> StreamingResponse:
|
||||
await transact(session, egress.require, egress.OLLAMA_PULL)
|
||||
await transact(
|
||||
session,
|
||||
egress.require,
|
||||
egress.OLLAMA_PULL,
|
||||
egress.ollama_pull_host(payload.name),
|
||||
)
|
||||
lock = service.install_lock(store.name)
|
||||
if lock.locked():
|
||||
raise HTTPException(
|
||||
|
|
|
|||
|
|
@ -1,10 +1,15 @@
|
|||
"""Nothing leaves the machine until the user allows that destination."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from httpx import AsyncClient
|
||||
from sqlalchemy import Engine
|
||||
|
||||
from modules.llm.models import ProviderConnection
|
||||
from modules.llm.recommendations.dependencies import get_catalog_service
|
||||
from shared.config import get_llm_settings
|
||||
from shared.db import create_session_factory
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
|
@ -95,3 +100,71 @@ async def test_unknown_destination_is_rejected(client: AsyncClient) -> None:
|
|||
"""Only destinations the app contacts can be toggled."""
|
||||
reply = await client.put("/egress/keygen", json={"enabled": True})
|
||||
assert reply.status_code == 422
|
||||
|
||||
|
||||
def _configure_llmfit(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
system = {"system": {"available_ram_gb": 16, "total_ram_gb": 16}}
|
||||
fit = {
|
||||
"models": [
|
||||
{
|
||||
"name": "Qwen/Qwen3-8B",
|
||||
"provider": "Qwen",
|
||||
"parameter_count": "8B",
|
||||
"use_case": "chat",
|
||||
"fit_level": "good",
|
||||
"score": 85,
|
||||
"runtime": "llamacpp",
|
||||
"run_mode": "gpu",
|
||||
"best_quant": "Q4_K_M",
|
||||
"memory_required_gb": 6,
|
||||
"memory_available_gb": 16,
|
||||
"disk_size_gb": 5.2,
|
||||
"effective_context_length": 8192,
|
||||
"capability_ids": ["tool_use"],
|
||||
"ollama_name": "qwen3:8b",
|
||||
"gguf_sources": [],
|
||||
}
|
||||
]
|
||||
}
|
||||
executable = tmp_path / "llmfit"
|
||||
executable.write_text(
|
||||
"#!/usr/bin/env python3\n"
|
||||
"import sys\n"
|
||||
f"system = {json.dumps(system)!r}\n"
|
||||
f"fit = {json.dumps(fit)!r}\n"
|
||||
'print("llmfit 1.1.11" if "--version" in sys.argv '
|
||||
'else system if "system" in sys.argv else fit)\n'
|
||||
)
|
||||
executable.chmod(0o755)
|
||||
monkeypatch.setattr(get_llm_settings(), "llmfit_path", executable)
|
||||
get_catalog_service.cache_clear()
|
||||
|
||||
|
||||
async def test_catalog_install_is_refused_until_allowed(
|
||||
client: AsyncClient,
|
||||
ollama_server: str,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""The catalog's own install route is gated the same as a manual pull."""
|
||||
_configure_llmfit(tmp_path, monkeypatch)
|
||||
catalog = (await client.get("/llm/catalog?refresh=true")).json()
|
||||
catalog_id = next(
|
||||
row["catalog_id"] for row in catalog["curated"] if row["fit"] != "unknown"
|
||||
)
|
||||
|
||||
refused = await client.post(
|
||||
"/llm/install", json={"catalog_id": catalog_id, "select": False}
|
||||
)
|
||||
assert refused.status_code == 403
|
||||
assert refused.json()["detail"]["destination"] == "ollama_pull"
|
||||
assert (await _destinations(client))["ollama_pull"]["last_call_at"] is None
|
||||
|
||||
await client.put("/egress/ollama_pull", json={"enabled": True})
|
||||
async with client.stream(
|
||||
"POST", "/llm/install", json={"catalog_id": catalog_id, "select": False}
|
||||
) as reply:
|
||||
assert reply.status_code == 200
|
||||
await reply.aread()
|
||||
assert (await _destinations(client))["ollama_pull"]["last_call_at"] is not None
|
||||
get_catalog_service.cache_clear()
|
||||
|
|
|
|||
|
|
@ -485,6 +485,7 @@ async def test_ranked_catalog_installs_and_selects_in_one_stream(
|
|||
assert [row["canonical_id"] for row in scanned_curated] == ["Qwen/Qwen3-8B"]
|
||||
catalog_id = scanned_curated[0]["catalog_id"]
|
||||
|
||||
await client.put("/egress/ollama_pull", json={"enabled": True})
|
||||
events = []
|
||||
async with client.stream(
|
||||
"POST",
|
||||
|
|
|
|||
15
surfsense_local/backend/tests/unit/egress/test_service.py
Normal file
15
surfsense_local/backend/tests/unit/egress/test_service.py
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
import pytest
|
||||
|
||||
from modules.egress import service as egress
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
|
||||
def test_a_library_pull_reports_the_ollama_registry() -> None:
|
||||
"""Ollama's own model library is served from its registry."""
|
||||
assert egress.ollama_pull_host("qwen3:8b") == "registry.ollama.ai"
|
||||
|
||||
|
||||
def test_an_hf_co_fallback_pull_reports_hugging_face() -> None:
|
||||
"""Ollama fetches `hf.co/<repo>` pulls straight from Hugging Face, not its own registry."""
|
||||
assert egress.ollama_pull_host("hf.co/Qwen/Qwen3-8B-GGUF:Q4_K_M") == "huggingface.co"
|
||||
Loading…
Add table
Add a link
Reference in a new issue