prometheusnn 0.2.0__tar.gz → 0.2.2__tar.gz

{prometheusnn-0.2.0 → prometheusnn-0.2.2}/CHANGELOG.md

@@ -5,6 +5,45 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.2.2] - 2026-10-30
+
+### Security
+- **Novelty data is loaded without pickle.** `novelty_centroids.npz` was loaded
+  with `allow_pickle=True`, so a crafted model directory could run code in the
+  server when any tool loaded that model. In 0.2.0 and earlier this was
+  reachable through the `model_name` path traversal fixed in 0.2.1; in 0.2.1 it
+  still applied to a malicious model placed in the models directory. Both
+  novelty files now load with `allow_pickle=False`. Models saved by
+  PrometheusNN contain only numeric and string arrays and load unchanged; a
+  file holding pickled objects is refused with an error asking for a retrain.
+
+## [0.2.1] - 2026-10-30
+
+### Security
+- **Model-name path traversal.** `delete_model`, `describe_model`,
+  `submit_feedback` and every other tool that takes a `model_name` now accept
+  only a single path segment (`[A-Za-z0-9._-]`, no `..`), and every resolved
+  path must stay under its base directory. In 0.2.0, `delete_model("..")`
+  deleted all of `PROMETHEUS_HOME`, an absolute name could delete any
+  directory the process could write, and `submit_feedback` could append to a
+  file outside the feedback directory. Upgrade if the MCP server or the UI
+  server is reachable by anything you do not fully trust.
+- **Unauthenticated UI WebSocket.** `prometheusnn-ui` now requires a token
+  (`?token=` on the WebSocket URL; set with `--token` or
+  `PROMETHEUS_UI_TOKEN`, otherwise generated and printed at startup) and
+  rejects browser handshakes from origins outside an allowlist (the hosted
+  demo and `localhost:3000` by default; add more with `--allowed-origin`).
+  In 0.2.0 any web page open in the same browser could drive every tool.
+- `delete_model` and `export_model` are withheld from the UI unless the
+  server is started with `--allow-destructive`.
+
+### Changed
+- `delete_model` refuses a directory that has no `manifest.json`.
+- `export_model` writes only under the exports directory
+  (`PROMETHEUS_EXPORT_DIR`, default `<PROMETHEUS_HOME>/exports`). A relative
+  `output_path` is placed there; an absolute path outside it is refused.
+  **Breaking** for callers that exported to arbitrary absolute paths.
+
 ## [0.2.0] - 2026-04-30
 
 ### Added

{prometheusnn-0.2.0 → prometheusnn-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: prometheusnn
-Version: 0.2.0
+Version: 0.2.2
 Summary: Frictionless hierarchical text classification — an MCP server that guides you from CSV to trained classifier
 Project-URL: Homepage, https://github.com/mfbaig35r/prometheusnn
 Project-URL: Repository, https://github.com/mfbaig35r/prometheusnn

{prometheusnn-0.2.0 → prometheusnn-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "prometheusnn"
-version = "0.2.0"
+version = "0.2.2"
 description = "Frictionless hierarchical text classification — an MCP server that guides you from CSV to trained classifier"
 readme = "README.md"
 requires-python = ">=3.11"

prometheusnn-0.2.2/src/prometheus/core/config.py

@@ -0,0 +1,106 @@
+"""Configuration and directory management."""
+
+import os
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from prometheus.core.mcp_safety import ToolValidationError
+
+# A model name is one path segment: letters, digits, dot, underscore, hyphen; no leading dot.
+_MODEL_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$")
+
+
+def validate_model_name(model_name: str) -> str:
+    """Reject anything that is not a single, plain directory name (no '/', '..', absolute paths)."""
+    if not isinstance(model_name, str) or not _MODEL_NAME.match(model_name) or ".." in model_name:
+        raise ToolValidationError(
+            f"Invalid model name {model_name!r}: use letters, digits, '.', '_' or '-' "
+            "(max 128 characters, no path separators)",
+            field="model_name",
+        )
+    return model_name
+
+
+def _contained(base: Path, child: Path) -> Path:
+    """Resolve child and require it to stay inside base (defence in depth after validation)."""
+    base_r, child_r = base.resolve(), child.resolve()
+    if child_r != base_r and base_r not in child_r.parents:
+        raise ToolValidationError(f"Path escapes {base}", field="model_name")
+    return child_r
+
+
+@dataclass
+class PrometheusConfig:
+    """Global configuration for Prometheus."""
+
+    base_dir: Path = field(
+        default_factory=lambda: Path(os.environ.get("PROMETHEUS_HOME", Path.home() / ".prometheus"))
+    )
+    default_embedding_model: str = field(
+        default_factory=lambda: os.environ.get("PROMETHEUS_EMBEDDING_MODEL", "all-MiniLM-L6-v2")
+    )
+
+    def __post_init__(self):
+        self.base_dir = Path(self.base_dir)
+
+    @property
+    def models_dir(self) -> Path:
+        return self.base_dir / "models"
+
+    @property
+    def feedback_dir(self) -> Path:
+        return self.base_dir / "feedback"
+
+    @property
+    def logs_dir(self) -> Path:
+        return self.base_dir / "logs"
+
+    @property
+    def exports_dir(self) -> Path:
+        override = os.environ.get("PROMETHEUS_EXPORT_DIR")
+        return Path(override) if override else self.base_dir / "exports"
+
+    def ensure_dirs(self):
+        """Create all required directories."""
+        self.models_dir.mkdir(parents=True, exist_ok=True)
+        self.feedback_dir.mkdir(parents=True, exist_ok=True)
+        self.logs_dir.mkdir(parents=True, exist_ok=True)
+
+    def model_path(self, model_name: str) -> Path:
+        """Directory for a model; always inside models_dir."""
+        return _contained(self.models_dir, self.models_dir / validate_model_name(model_name))
+
+    def feedback_path(self, model_name: str) -> Path:
+        """Feedback directory for a model; always inside feedback_dir."""
+        return _contained(self.feedback_dir, self.feedback_dir / validate_model_name(model_name))
+
+    def export_path(self, output_path: str) -> Path:
+        """Destination for export_model; always inside exports_dir.
+
+        A relative path is taken under exports_dir; an absolute path must already point inside
+        it. Set PROMETHEUS_EXPORT_DIR to export somewhere else.
+        """
+        if not isinstance(output_path, str) or not output_path.strip():
+            raise ToolValidationError("output_path is required", field="output_path")
+        candidate = Path(output_path).expanduser()
+        if not candidate.is_absolute():
+            candidate = self.exports_dir / candidate
+        try:
+            return _contained(self.exports_dir, candidate)
+        except ToolValidationError as exc:
+            raise ToolValidationError(
+                f"output_path must be inside the export directory {self.exports_dir} "
+                "(set PROMETHEUS_EXPORT_DIR to change it)",
+                field="output_path",
+            ) from exc
+
+    def list_models(self) -> list[str]:
+        """List all saved model names."""
+        if not self.models_dir.exists():
+            return []
+        return sorted(
+            d.name
+            for d in self.models_dir.iterdir()
+            if d.is_dir() and (d / "manifest.json").exists()
+        )

{prometheusnn-0.2.0 → prometheusnn-0.2.2}/src/prometheus/inference/novelty.py

@@ -116,17 +116,26 @@ class NoveltyDetector:
 
     def load(self, path: Path):
         path = Path(path)
-        data = np.load(path / "novelty_centroids.npz", allow_pickle=True)
-        class_names = data["class_names"]
-        mean_dists = data["mean_dists"]
-        std_dists = data["std_dists"]
-        p95_dists = data["p95_dists"]
-        counts = data["counts"]
+        # Never unpickle: a model directory is data, and a crafted .npz with object arrays would
+        # run code on load. save() writes only numeric and string arrays, which need no pickle.
+        data = np.load(path / "novelty_centroids.npz", allow_pickle=False)
+        try:
+            class_names = data["class_names"]
+            mean_dists = data["mean_dists"]
+            std_dists = data["std_dists"]
+            p95_dists = data["p95_dists"]
+            counts = data["counts"]
+            centroids = [data[f"centroid_{i}"] for i in range(len(class_names))]
+        except ValueError as exc:
+            raise ValueError(
+                f"{path / 'novelty_centroids.npz'} contains pickled objects; refusing to load it. "
+                "Retrain the model to regenerate it."
+            ) from exc
 
         self.centroids = {}
         for i, name in enumerate(class_names):
             self.centroids[str(name)] = ClassCentroid(
-                centroid=data[f"centroid_{i}"],
+                centroid=centroids[i],
                 mean_dist=float(mean_dists[i]),
                 std_dist=max(float(std_dists[i]), 0.1),
                 p95_dist=float(p95_dists[i]),
@@ -135,7 +144,7 @@ class NoveltyDetector:
 
         emb_path = path / "novelty_embeddings.npy"
         if emb_path.exists():
-            self.training_embeddings = np.load(emb_path)
+            self.training_embeddings = np.load(emb_path, allow_pickle=False)
 
     def get_centroids_data(self) -> dict:
         """Return data dict suitable for np.savez."""

{prometheusnn-0.2.0 → prometheusnn-0.2.2}/src/prometheus/server/tools/evaluation.py

@@ -242,7 +242,7 @@ def submit_feedback(
         prediction_id: Optional prediction ID to link feedback to
     """
     logger.info("tool=submit_feedback model=%s prediction_id=%s", model_name, prediction_id)
-    feedback_dir = server._config.feedback_dir / model_name
+    feedback_dir = server._config.feedback_path(model_name)
     feedback_dir.mkdir(parents=True, exist_ok=True)
 
     entry = {

{prometheusnn-0.2.0 → prometheusnn-0.2.2}/src/prometheus/server/tools/management.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 
 import json
 import shutil
-from pathlib import Path
 
 from prometheus import server
 from prometheus.core.mcp_safety import ToolValidationError, mcp_tool_safe
@@ -86,6 +85,9 @@ def delete_model(model_name: str) -> dict:
     model_path = server._config.model_path(model_name)
     if not model_path.exists():
         raise ToolValidationError(f"Model '{model_name}' not found")
+    # Only ever delete a directory that is a saved model.
+    if not (model_path / "manifest.json").is_file():
+        raise ToolValidationError(f"'{model_name}' is not a saved model (no manifest.json)")
 
     delete_model_from_disk(model_path)
     server._predictor_cache.pop(model_name, None)
@@ -99,21 +101,24 @@ def export_model(
     model_name: str,
     output_path: str,
 ) -> dict:
-    """Copy model directory to a new location.
+    """Copy model directory into the export directory.
 
     Args:
         model_name: Name of the model to export
-        output_path: Destination directory
+        output_path: Destination directory, relative to the export directory
+            (PROMETHEUS_HOME/exports by default, or PROMETHEUS_EXPORT_DIR). An absolute path
+            must point inside it.
     """
     logger.info("tool=export_model model=%s dest=%s", model_name, output_path)
     model_path = server._config.model_path(model_name)
     if not model_path.exists():
         raise ToolValidationError(f"Model '{model_name}' not found")
 
-    dest = Path(output_path)
+    dest = server._config.export_path(output_path)
     if dest.exists():
         raise ToolValidationError(f"Destination already exists: {output_path}")
 
+    dest.parent.mkdir(parents=True, exist_ok=True)
     shutil.copytree(model_path, dest)
     return success_response({"exported_to": str(dest)})
 

{prometheusnn-0.2.0 → prometheusnn-0.2.2}/src/prometheus/ui_server.py

@@ -6,6 +6,14 @@ https://prometheusnn.cognitionkernel.agency — toggle to "connected mode"
 and point it at this server (default: ws://localhost:8000/ws/prometheus)
 to drive the same UI against your own data, on your own hardware.
 
+Security. Browsers let any web page open a WebSocket to localhost, so the server:
+
+- requires a per-launch token, passed as ``?token=...`` in the URL (printed at startup, or set
+  with ``--token`` / ``PROMETHEUS_UI_TOKEN``); paste the printed URL into connected mode;
+- accepts browser connections only from allowed origins (the companion site and local dev by
+  default; ``--allowed-origin`` to change);
+- does not offer ``delete_model`` or ``export_model`` unless started with ``--allow-destructive``.
+
 Protocol — JSON over WebSocket. Atomic request/response per tool call.
 
 Client → Server:
@@ -25,9 +33,12 @@ from __future__ import annotations
 
 import argparse
 import asyncio
+import hmac
 import inspect
 import json
 import logging
+import os
+import secrets
 from collections.abc import Callable
 from typing import Any
 
@@ -49,6 +60,58 @@ def _build_tool_registry() -> dict[str, Callable[..., dict]]:
     return {name: tool.fn for name, tool in tool_manager._tools.items()}
 
 
+# ─── Connection security ─────────────────────────────────────────────────────
+
+DEFAULT_ALLOWED_ORIGINS = (
+    "https://prometheusnn.cognitionkernel.agency",
+    "http://localhost:3000",
+    "http://127.0.0.1:3000",
+)
+# Tools that delete or write outside the model store; off over the UI unless --allow-destructive.
+DESTRUCTIVE_TOOLS = frozenset({"delete_model", "export_model"})
+
+_SECURITY: dict[str, Any] = {
+    "token": None,
+    "allowed_origins": frozenset(DEFAULT_ALLOWED_ORIGINS),
+    "allow_destructive": False,
+}
+
+
+def configure(
+    *,
+    token: str | None,
+    allowed_origins: tuple[str, ...] | list[str] = DEFAULT_ALLOWED_ORIGINS,
+    allow_destructive: bool = False,
+) -> None:
+    """Set connection security. main() calls this; tests call it directly."""
+    _SECURITY["token"] = token
+    _SECURITY["allowed_origins"] = frozenset(o.rstrip("/") for o in allowed_origins)
+    _SECURITY["allow_destructive"] = allow_destructive
+
+
+def check_handshake(origin: str | None, token: str | None) -> str | None:
+    """None if the connection may proceed, otherwise the reason it is refused.
+
+    Browsers always send Origin on a WebSocket handshake, so a present Origin must be allowed;
+    a missing Origin means a non-browser client, which still needs the token.
+    """
+    expected = _SECURITY["token"]
+    if not expected:
+        return "server has no token configured"
+    if origin is not None and origin.rstrip("/") not in _SECURITY["allowed_origins"]:
+        return f"origin not allowed: {origin}"
+    if not token or not hmac.compare_digest(token, expected):
+        return "missing or invalid token"
+    return None
+
+
+def _ui_tools() -> dict[str, Callable[..., dict]]:
+    tools = _get_tools()
+    if _SECURITY["allow_destructive"]:
+        return tools
+    return {n: f for n, f in tools.items() if n not in DESTRUCTIVE_TOOLS}
+
+
 _TOOLS_CACHE: dict[str, Callable[..., dict]] | None = None
 
 
@@ -90,7 +153,7 @@ def _serialize_tool_metadata(tool_fn: Callable[..., dict]) -> dict[str, Any]:
 
 async def _handle_message(ws, msg: dict[str, Any]) -> None:
     msg_type = msg.get("type")
-    tools = _get_tools()
+    tools = _ui_tools()
 
     if msg_type == "ping":
         await ws.send_json({"type": "pong"})
@@ -111,12 +174,22 @@ async def _handle_message(ws, msg: dict[str, Any]) -> None:
         args = msg.get("args", {}) or {}
 
         if tool_name not in tools:
+            error = (
+                f"{tool_name} is not available over the UI server (start it with --allow-destructive)"
+                if tool_name in DESTRUCTIVE_TOOLS
+                else f"unknown tool: {tool_name}"
+            )
+            await ws.send_json(
+                {"type": "tool_error", "id": call_id, "tool": tool_name, "error": error}
+            )
+            return
+        if not isinstance(args, dict):
             await ws.send_json(
                 {
                     "type": "tool_error",
                     "id": call_id,
                     "tool": tool_name,
-                    "error": f"unknown tool: {tool_name}",
+                    "error": "args must be an object",
                 }
             )
             return
@@ -163,7 +236,7 @@ async def _handle_message(ws, msg: dict[str, Any]) -> None:
 try:
     from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 
-    app: FastAPI | None = FastAPI(title="prometheusnn-ui", version="0.2.0")
+    app: FastAPI | None = FastAPI(title="prometheusnn-ui", version="0.2.2")
 
     @app.get("/health")
     def _health() -> dict:
@@ -171,6 +244,13 @@ try:
 
     @app.websocket("/ws/prometheus")
     async def _ws_endpoint(websocket: WebSocket) -> None:
+        refused = check_handshake(
+            websocket.headers.get("origin"), websocket.query_params.get("token")
+        )
+        if refused is not None:
+            logger.warning("refused WebSocket connection: %s", refused)
+            await websocket.close(code=1008)  # before accept(): the handshake fails
+            return
         await websocket.accept()
         peer = f"{websocket.client.host}:{websocket.client.port}" if websocket.client else "unknown"
         logger.info("client connected: %s", peer)
@@ -191,12 +271,21 @@ except ImportError:
 
 
 def create_app():
-    """Factory used by uvicorn --factory. Returns the module-level app."""
+    """Factory used by uvicorn --factory. Returns the module-level app.
+
+    Outside main() (e.g. ``uvicorn --factory``), the token comes from PROMETHEUS_UI_TOKEN or is
+    generated and printed, so the server never runs without one.
+    """
     if app is None:
         raise SystemExit(
             "prometheusnn-ui requires the [ui] extras. Install with:\n"
             "    pip install prometheusnn[ui]"
         )
+    if not _SECURITY["token"]:
+        token = os.environ.get("PROMETHEUS_UI_TOKEN") or secrets.token_urlsafe(24)
+        configure(token=token)
+        if "PROMETHEUS_UI_TOKEN" not in os.environ:
+            print(f"\n  prometheusnn-ui token: {token}  (append ?token=... to the URL)\n")
     return app
 
 
@@ -216,6 +305,23 @@ def main() -> None:
         default="info",
         choices=["debug", "info", "warning", "error"],
     )
+    parser.add_argument(
+        "--token",
+        default=os.environ.get("PROMETHEUS_UI_TOKEN"),
+        help="Connection token (default: PROMETHEUS_UI_TOKEN, or a new random one per launch)",
+    )
+    parser.add_argument(
+        "--allowed-origin",
+        action="append",
+        dest="allowed_origins",
+        help="Browser origin allowed to connect; repeatable (default: the companion site and "
+        "localhost:3000)",
+    )
+    parser.add_argument(
+        "--allow-destructive",
+        action="store_true",
+        help="Also offer delete_model and export_model over the UI",
+    )
     args = parser.parse_args()
 
     try:
@@ -227,12 +333,21 @@ def main() -> None:
         ) from e
 
     logging.basicConfig(level=args.log_level.upper())
+    token = args.token or secrets.token_urlsafe(24)
+    configure(
+        token=token,
+        allowed_origins=args.allowed_origins or DEFAULT_ALLOWED_ORIGINS,
+        allow_destructive=args.allow_destructive,
+    )
+    if args.host not in ("127.0.0.1", "localhost", "::1"):
+        logger.warning("binding to %s exposes the server beyond this machine", args.host)
     # Warm the tool registry so the count is logged before serving
     tools = _get_tools()
     logger.info("Registered %d tools", len(tools))
     print(
         f"\n  prometheusnn-ui listening on ws://{args.host}:{args.port}/ws/prometheus\n"
-        f"  Open https://prometheusnn.cognitionkernel.agency and toggle 'Connected mode'\n"
+        f"  Open https://prometheusnn.cognitionkernel.agency, toggle 'Connected mode', and use:\n"
+        f"    ws://{args.host}:{args.port}/ws/prometheus?token={token}\n"
     )
     uvicorn.run(
         create_app(),

prometheusnn-0.2.2/tests/test_novelty.py

@@ -0,0 +1,100 @@
+"""Tests for novelty detection."""
+
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from prometheus.inference.novelty import NoveltyDetector
+
+
+def test_fit_and_score():
+    rng = np.random.default_rng(42)
+    n = 50
+    embs = rng.standard_normal((n, 16)).astype(np.float32)
+    labels = np.array([0] * 25 + [1] * 25)
+    classes = ["classA", "classB"]
+
+    det = NoveltyDetector(z_score_threshold=2.5)
+    det.fit(embs, labels, classes)
+
+    assert len(det.centroids) == 2
+
+    # Score an in-distribution sample
+    score = det.score(embs[0], "classA", confidence=0.9)
+    assert not score.is_novel
+
+    # Score a far-out sample
+    outlier = np.ones(16, dtype=np.float32) * 100
+    score = det.score(outlier, "classA", confidence=0.3)
+    assert score.centroid_z_score > 2.5
+
+
+def test_save_load_roundtrip():
+    rng = np.random.default_rng(42)
+    embs = rng.standard_normal((30, 8)).astype(np.float32)
+    labels = np.array([0] * 15 + [1] * 15)
+    classes = ["X", "Y"]
+
+    det = NoveltyDetector()
+    det.fit(embs, labels, classes)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        det.save(Path(tmpdir))
+
+        det2 = NoveltyDetector()
+        det2.load(Path(tmpdir))
+
+        assert set(det2.centroids.keys()) == {"X", "Y"}
+        assert det2.training_embeddings is not None
+        assert det2.training_embeddings.shape == (30, 8)
+
+
+class _Payload:
+    """Unpickling this creates a file: proof that loading ran attacker code."""
+
+    def __init__(self, marker: Path):
+        self.marker = str(marker)
+
+    def __reduce__(self):
+        return (open, (self.marker, "w"))
+
+
+def test_load_never_unpickles_centroids(tmp_path):
+    marker = tmp_path / "executed"
+    np.savez(
+        tmp_path / "novelty_centroids.npz",
+        centroid_0=np.zeros(4),
+        class_names=np.array([_Payload(marker)], dtype=object),
+        mean_dists=[1.0],
+        std_dists=[1.0],
+        p95_dists=[1.0],
+        counts=[1],
+    )
+    with pytest.raises(ValueError, match="pickled objects"):
+        NoveltyDetector().load(tmp_path)
+    assert not marker.exists()
+
+
+def test_load_never_unpickles_embeddings(tmp_path):
+    marker = tmp_path / "executed"
+    det = NoveltyDetector()
+    det.fit(np.eye(4, dtype=np.float32), np.array([0, 0, 1, 1]), ["a", "b"])
+    det.save(tmp_path)
+    np.save(tmp_path / "novelty_embeddings.npy", np.array([_Payload(marker)], dtype=object))
+    with pytest.raises(ValueError):
+        NoveltyDetector().load(tmp_path)
+    assert not marker.exists()
+
+
+def test_saved_files_need_no_pickle(tmp_path):
+    rng = np.random.default_rng(0)
+    det = NoveltyDetector()
+    det.fit(rng.standard_normal((20, 4)), np.array([0] * 10 + [1] * 10), ["Café", "naïve/class"])
+    det.save(tmp_path)
+    with np.load(tmp_path / "novelty_centroids.npz", allow_pickle=False) as data:
+        assert all(data[k].dtype != object for k in data.files)
+    det2 = NoveltyDetector()
+    det2.load(tmp_path)
+    assert set(det2.centroids) == {"Café", "naïve/class"}

{prometheusnn-0.2.0 → prometheusnn-0.2.2}/tests/test_server.py

@@ -267,11 +267,92 @@ class TestDeleteModel:
 
 
 class TestExportModel:
-    def test_export(self, trained_model, tmp_path):
-        dest = str(tmp_path / "exported_model")
-        result = server.export_model(trained_model, dest)
+    def test_export(self, trained_model, tmp_config, monkeypatch):
+        monkeypatch.delenv("PROMETHEUS_EXPORT_DIR", raising=False)
+        result = server.export_model(trained_model, "exported_model")
         assert result["status"] == "success"
-        assert Path(dest).exists()
+        assert (tmp_config.exports_dir / "exported_model" / "manifest.json").exists()
+
+    def test_export_outside_export_dir_is_refused(self, trained_model, tmp_path, monkeypatch):
+        monkeypatch.delenv("PROMETHEUS_EXPORT_DIR", raising=False)
+        for dest in (str(tmp_path / "elsewhere"), "../../escape", "/tmp/prometheus-escape"):
+            result = server.export_model(trained_model, dest)
+            assert result["status"] == "validation_error", dest
+        assert not (tmp_path / "elsewhere").exists()
+
+    def test_export_dir_override(self, trained_model, tmp_path, monkeypatch):
+        target = tmp_path / "my-exports"
+        monkeypatch.setenv("PROMETHEUS_EXPORT_DIR", str(target))
+        result = server.export_model(trained_model, str(target / "copy"))
+        assert result["status"] == "success"
+        assert (target / "copy" / "manifest.json").exists()
+
+
+# ─── Model-name path traversal (security) ───
+
+BAD_NAMES = [
+    # Never list a real system path here: against vulnerable code these names reach rmtree.
+    "../escape",
+    "../../nonexistent-parent",
+    "/nonexistent-prometheus-test-path",
+    "/tmp/nonexistent-prometheus-test-dir",
+    "a/b",
+    "..",
+    ".hidden",
+    "",
+    "x" * 200,
+    "name\\with\\backslash",
+    "name\x00null",
+]
+
+
+class TestModelNameTraversal:
+    @pytest.mark.parametrize("name", BAD_NAMES)
+    def test_delete_refuses_bad_names(self, name, tmp_path):
+        victim = tmp_path / "victim"
+        victim.mkdir()
+        (victim / "keep.txt").write_text("keep")
+        result = server.delete_model(name)
+        assert result["status"] == "validation_error"
+        assert (victim / "keep.txt").exists()
+
+    def test_delete_absolute_path_to_real_dir_is_refused(self, tmp_path):
+        victim = tmp_path / "important"
+        victim.mkdir()
+        (victim / "data.txt").write_text("precious")
+        result = server.delete_model(str(victim))
+        assert result["status"] == "validation_error"
+        assert (victim / "data.txt").read_text() == "precious"
+
+    def test_delete_refuses_a_directory_without_manifest(self, tmp_config):
+        stray = tmp_config.models_dir / "not_a_model"
+        stray.mkdir()
+        (stray / "file.txt").write_text("x")
+        result = server.delete_model("not_a_model")
+        assert result["status"] == "validation_error"
+        assert stray.exists()
+
+    @pytest.mark.parametrize("name", ["../outside", "/tmp/fb-outside"])
+    def test_feedback_cannot_write_outside_feedback_dir(self, name, tmp_config, tmp_path):
+        result = server.submit_feedback(name, "some text", {"level_1": "x"})
+        assert result["status"] == "validation_error"
+        assert not (tmp_config.feedback_dir.parent / "outside").exists()
+
+    @pytest.mark.parametrize("tool", ["describe_model", "predict", "evaluate_model"])
+    def test_other_tools_validate_names(self, tool):
+        fn = getattr(server, tool)
+        args = {
+            "describe_model": ("../x",),
+            "predict": ("../x", "text"),
+            "evaluate_model": ("../x",),
+        }[tool]
+        result = fn(*args)
+        assert result["status"] in ("validation_error", "error"), result
+
+    @pytest.mark.parametrize("name", ["model-1", "my_model.v2", "A1", "x" * 128])
+    def test_good_names_resolve_inside_models_dir(self, name, tmp_config):
+        path = tmp_config.model_path(name)
+        assert path.parent == tmp_config.models_dir.resolve()
 
 
 # ─── Code Mapping ───

prometheusnn-0.2.2/tests/test_ui_server.py

@@ -0,0 +1,101 @@
+"""WebSocket UI server: origin allowlist, per-launch token, destructive tools withheld."""
+
+from __future__ import annotations
+
+import pytest
+
+fastapi = pytest.importorskip("fastapi")
+from fastapi.testclient import TestClient  # noqa: E402
+from starlette.websockets import WebSocketDisconnect  # noqa: E402
+
+from prometheus import ui_server  # noqa: E402
+
+TOKEN = "test-token-123"
+GOOD = "https://prometheusnn.cognitionkernel.agency"
+
+
+@pytest.fixture(autouse=True)
+def secured():
+    ui_server.configure(token=TOKEN)
+    yield
+    ui_server.configure(token=None)
+
+
+@pytest.fixture
+def client():
+    return TestClient(ui_server.app)
+
+
+def _connect(client, origin=GOOD, token=TOKEN):
+    url = "/ws/prometheus" + (f"?token={token}" if token is not None else "")
+    headers = {"origin": origin} if origin is not None else {}
+    return client.websocket_connect(url, headers=headers)
+
+
+def test_allowed_origin_with_token_connects(client):
+    with _connect(client) as ws:
+        ws.send_json({"type": "ping"})
+        assert ws.receive_json() == {"type": "pong"}
+
+
+@pytest.mark.parametrize("origin", ["https://evil.example", "http://localhost:9999", "null"])
+def test_foreign_origin_is_refused(client, origin):
+    with pytest.raises(WebSocketDisconnect) as exc, _connect(client, origin=origin):
+        pass
+    assert exc.value.code == 1008
+
+
+@pytest.mark.parametrize("token", [None, "", "wrong"])
+def test_missing_or_wrong_token_is_refused(client, token):
+    with pytest.raises(WebSocketDisconnect), _connect(client, token=token):
+        pass
+
+
+def test_non_browser_client_still_needs_the_token(client):
+    with _connect(client, origin=None) as ws:
+        ws.send_json({"type": "ping"})
+        assert ws.receive_json()["type"] == "pong"
+    with pytest.raises(WebSocketDisconnect), _connect(client, origin=None, token="nope"):
+        pass
+
+
+def test_no_token_configured_refuses_everything(client):
+    ui_server.configure(token=None)
+    with pytest.raises(WebSocketDisconnect), _connect(client):
+        pass
+
+
+def test_destructive_tools_are_withheld_by_default(client):
+    with _connect(client) as ws:
+        ws.send_json({"type": "list_tools"})
+        names = {t["name"] for t in ws.receive_json()["tools"]}
+        assert not names & ui_server.DESTRUCTIVE_TOOLS
+        assert "start_here" in names
+        ws.send_json(
+            {"type": "tool_call", "id": "1", "tool": "delete_model", "args": {"model_name": "x"}}
+        )
+        msg = ws.receive_json()
+        assert msg["type"] == "tool_error" and "--allow-destructive" in msg["error"]
+
+
+def test_allow_destructive_opt_in(client):
+    ui_server.configure(token=TOKEN, allow_destructive=True)
+    with _connect(client) as ws:
+        ws.send_json({"type": "list_tools"})
+        names = {t["name"] for t in ws.receive_json()["tools"]}
+        assert names >= ui_server.DESTRUCTIVE_TOOLS
+
+
+def test_non_object_args_are_rejected(client):
+    with _connect(client) as ws:
+        ws.send_json({"type": "tool_call", "id": "2", "tool": "list_models", "args": ["a"]})
+        assert ws.receive_json()["error"] == "args must be an object"
+
+
+def test_custom_origin_list(client):
+    ui_server.configure(token=TOKEN, allowed_origins=["http://localhost:5173"])
+    with _connect(client, origin="http://localhost:5173") as ws:
+        ws.send_json({"type": "ping"})
+        assert ws.receive_json()["type"] == "pong"
+    with pytest.raises(WebSocketDisconnect), _connect(client, origin=GOOD):
+        pass

{prometheusnn-0.2.0 → prometheusnn-0.2.2}/uv.lock

@@ -907,7 +907,7 @@ wheels = [
 
 [[package]]
 name = "prometheusnn"
-version = "0.1.0"
+version = "0.2.2"
 source = { editable = "." }
 dependencies = [
     { name = "mcp", extra = ["cli"] },

prometheusnn-0.2.0/src/prometheus/core/config.py

@@ -1,51 +0,0 @@
-"""Configuration and directory management."""
-
-import os
-from dataclasses import dataclass, field
-from pathlib import Path
-
-
-@dataclass
-class PrometheusConfig:
-    """Global configuration for Prometheus."""
-
-    base_dir: Path = field(
-        default_factory=lambda: Path(os.environ.get("PROMETHEUS_HOME", Path.home() / ".prometheus"))
-    )
-    default_embedding_model: str = field(
-        default_factory=lambda: os.environ.get("PROMETHEUS_EMBEDDING_MODEL", "all-MiniLM-L6-v2")
-    )
-
-    def __post_init__(self):
-        self.base_dir = Path(self.base_dir)
-
-    @property
-    def models_dir(self) -> Path:
-        return self.base_dir / "models"
-
-    @property
-    def feedback_dir(self) -> Path:
-        return self.base_dir / "feedback"
-
-    @property
-    def logs_dir(self) -> Path:
-        return self.base_dir / "logs"
-
-    def ensure_dirs(self):
-        """Create all required directories."""
-        self.models_dir.mkdir(parents=True, exist_ok=True)
-        self.feedback_dir.mkdir(parents=True, exist_ok=True)
-        self.logs_dir.mkdir(parents=True, exist_ok=True)
-
-    def model_path(self, model_name: str) -> Path:
-        return self.models_dir / model_name
-
-    def list_models(self) -> list[str]:
-        """List all saved model names."""
-        if not self.models_dir.exists():
-            return []
-        return sorted(
-            d.name
-            for d in self.models_dir.iterdir()
-            if d.is_dir() and (d / "manifest.json").exists()
-        )

prometheusnn-0.2.0/tests/test_novelty.py

@@ -1,50 +0,0 @@
-"""Tests for novelty detection."""
-
-import tempfile
-from pathlib import Path
-
-import numpy as np
-
-from prometheus.inference.novelty import NoveltyDetector
-
-
-def test_fit_and_score():
-    rng = np.random.default_rng(42)
-    n = 50
-    embs = rng.standard_normal((n, 16)).astype(np.float32)
-    labels = np.array([0] * 25 + [1] * 25)
-    classes = ["classA", "classB"]
-
-    det = NoveltyDetector(z_score_threshold=2.5)
-    det.fit(embs, labels, classes)
-
-    assert len(det.centroids) == 2
-
-    # Score an in-distribution sample
-    score = det.score(embs[0], "classA", confidence=0.9)
-    assert not score.is_novel
-
-    # Score a far-out sample
-    outlier = np.ones(16, dtype=np.float32) * 100
-    score = det.score(outlier, "classA", confidence=0.3)
-    assert score.centroid_z_score > 2.5
-
-
-def test_save_load_roundtrip():
-    rng = np.random.default_rng(42)
-    embs = rng.standard_normal((30, 8)).astype(np.float32)
-    labels = np.array([0] * 15 + [1] * 15)
-    classes = ["X", "Y"]
-
-    det = NoveltyDetector()
-    det.fit(embs, labels, classes)
-
-    with tempfile.TemporaryDirectory() as tmpdir:
-        det.save(Path(tmpdir))
-
-        det2 = NoveltyDetector()
-        det2.load(Path(tmpdir))
-
-        assert set(det2.centroids.keys()) == {"X", "Y"}
-        assert det2.training_embeddings is not None
-        assert det2.training_embeddings.shape == (30, 8)