Merge pull request #28 from hybridindie/fix/p1-quick-wins

fix: P1 quick-win findings from comprehensive review

Author: hybridindie

.github/workflows/docker.yml

@@ -9,7 +9,31 @@ on:
       - 'v*'
 
 jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5  # v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@e4db8464a088ece1b920f60402e813ea4de65b8f  # v4
+
+      - name: Set up Python
+        run: uv python install 3.12
+
+      - name: Install dependencies
+        run: uv sync
+
+      - name: Lint and type-check
+        run: |
+          uv run ruff check src/ tests/
+          uv run ruff format --check src/ tests/
+          uv run mypy src/comfyui_mcp/
+
+      - name: Run tests
+        run: uv run pytest -v
+
   docker:
+    needs: test
     runs-on: ubuntu-latest
     permissions:
       contents: read

Dockerfile

@@ -2,7 +2,7 @@ FROM python:3.12-slim
 
 WORKDIR /app
 
-RUN groupadd --system app && useradd --system --gid app app
+RUN groupadd --system app && useradd --system --gid app --create-home app
 
 RUN pip install --no-cache-dir uv
 

README.md

@@ -279,6 +279,8 @@ transport:
     port: 8080
 ```
 
+> **SSE transport security warning:** The SSE transport has no built-in authentication. Binding to `0.0.0.0` or any non-localhost address exposes all MCP tools (workflow execution, file uploads, queue management) to any network client without authentication. Only enable SSE behind a reverse proxy (e.g., nginx, Caddy) that provides authentication and TLS. The server logs a warning at startup when SSE is bound to a non-localhost address.
+
 ### Environment variables
 
 Environment variables override config file values:
@@ -411,6 +413,15 @@ security:
 
 **Tip:** Use `audit_dangerous_nodes` to identify dangerous nodes, run workflows in audit mode to see which nodes you use, then switch to enforce mode with that allowlist.
 
+### Choosing a security mode
+
+| Mode | Behavior | Use case |
+|------|----------|----------|
+| `audit` (default) | Logs warnings about dangerous nodes but **allows** all workflows to execute | Development, initial setup, discovering which nodes your workflows use |
+| `enforce` | **Blocks** any workflow containing a node not on `allowed_nodes` | Production, shared environments, security-sensitive deployments |
+
+Start in `audit` mode, review your audit log to build an allowlist, then switch to `enforce` mode for production use.
+
 ## Audit log
 
 All tool invocations are logged as JSON lines to `~/.comfyui-mcp/audit.log`:
@@ -425,6 +436,24 @@ grep '"warnings":\[' ~/.comfyui-mcp/audit.log | grep -v '"warnings":\[\]'
 
 Sensitive fields (`token`, `password`, `secret`, `api_key`, `authorization`) are automatically redacted from log entries.
 
+### Log rotation
+
+The audit log grows indefinitely. Set up log rotation to prevent disk exhaustion:
+
+```bash
+# /etc/logrotate.d/comfyui-mcp
+/home/user/.comfyui-mcp/audit.log {
+    weekly
+    rotate 12
+    compress
+    missingok
+    notifempty
+    copytruncate
+}
+```
+
+Use `copytruncate` because the server holds the file open. Adjust the path to match your `logging.audit_file` setting.
+
 ## Security
 
 ### Threat model
@@ -435,7 +464,6 @@ Sensitive fields (`token`, `password`, `secret`, `api_key`, `authorization`) are
 | Path traversal via file operations | High | Path sanitizer blocks `..`, null bytes, encoded attacks, absolute paths |
 | Denial of service via request flooding | Medium | Token-bucket rate limiter per tool category |
 | Credential leakage in logs | Medium | Automatic redaction of `token`, `password`, `secret`, `api_key`, `authorization` |
-| Information disclosure via API | Low | Dangerous endpoints (`/userdata`, `/free`, `/system_stats`) never proxied |
 | Information disclosure via API | Low | Dangerous endpoints (`/userdata`, `/free`) never proxied; `/system_stats` whitelist-filtered by `get_system_info` |
 | MITM on ComfyUI connection | Medium | Configurable TLS verification |
 

docker-compose.yml

@@ -1,5 +1,3 @@
-version: "3.8"
-
 services:
   comfyui-mcp:
     build: .
@@ -9,8 +7,8 @@ services:
       - COMFYUI_URL=${COMFYUI_URL:-http://comfyui:8188}
       - COMFYUI_SECURITY_MODE=${COMFYUI_SECURITY_MODE:-audit}
     volumes:
-      - ./config.yaml:/root/.comfyui-mcp/config.yaml:ro
-      - comfyui-mcp-data:/root/.comfyui-mcp/logs
+      - ./config.yaml:/home/app/.comfyui-mcp/config.yaml:ro
+      - comfyui-mcp-data:/home/app/.comfyui-mcp/logs
     restart: unless-stopped
 
 volumes:

src/comfyui_mcp/model_manager.py

@@ -3,11 +3,14 @@
 from __future__ import annotations
 
 import asyncio
+import time
 
 import httpx
 
 from comfyui_mcp.client import ComfyUIClient
 
+_NEGATIVE_TTL_SECONDS = 300  # Re-probe after 5 minutes on failure
+
 
 class ModelManagerUnavailableError(Exception):
     """Raised when Model Manager is not installed or unreachable."""
@@ -19,6 +22,10 @@ class ModelManagerDetector:
     Uses asyncio.Lock to prevent concurrent probes from racing.
     Calls GET /model-manager/models once — the response both confirms availability
     and provides the folder list in a single round-trip.
+
+    Positive results are cached permanently. Negative results expire after
+    ``_NEGATIVE_TTL_SECONDS`` so a late-starting Model Manager is picked up
+    without requiring a server restart.
     """
 
     _INSTALL_URL = "https://github.com/hayden-fr/ComfyUI-Model-Manager"
@@ -28,19 +35,30 @@ def __init__(self, client: ComfyUIClient) -> None:
         self._folders: list[str] | None = None
         self._checked = False
         self._available = False
+        self._last_failure: float = 0.0
         self._lock = asyncio.Lock()
 
+    def _should_reprobe(self) -> bool:
+        """Return True if the negative cache has expired."""
+        if self._available:
+            return False
+        if not self._checked:
+            return True
+        return (time.monotonic() - self._last_failure) >= _NEGATIVE_TTL_SECONDS
+
     async def _probe(self) -> None:
-        """Probe Model Manager once and cache the result."""
+        """Probe Model Manager and cache the result."""
         async with self._lock:
-            if self._checked:
+            if self._checked and not self._should_reprobe():
                 return
-            self._checked = True
             try:
                 self._folders = await self._client.get_model_manager_folders()
                 self._available = True
+                self._checked = True
             except (httpx.HTTPStatusError, httpx.RequestError):
                 self._available = False
+                self._checked = True
+                self._last_failure = time.monotonic()
 
     async def is_available(self) -> bool:
         """Check if Model Manager is installed. Caches the result."""

tests/test_client.py

@@ -220,6 +220,46 @@ async def test_no_retry_on_http_error(self, client):
         assert route.call_count == 1
 
 
+class TestPathInjectionPrevention:
+    async def test_prompt_id_path_traversal_rejected(self, client):
+        with pytest.raises(ValueError, match="Invalid prompt_id"):
+            await client.get_history_item("../../../free")
+
+    async def test_prompt_id_non_uuid_rejected(self, client):
+        with pytest.raises(ValueError, match="Invalid prompt_id"):
+            await client.get_history_item("not-a-uuid")
+
+    @respx.mock
+    async def test_prompt_id_valid_uuid_accepted(self, client):
+        route = respx.get(
+            "http://test-comfyui:8188/history/aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee"
+        ).mock(
+            return_value=httpx.Response(
+                200,
+                json={"aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee": {"outputs": {}}},
+            )
+        )
+        result = await client.get_history_item("aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee")
+        assert "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" in result
+        assert route.call_count == 1
+
+    async def test_delete_queue_path_traversal_rejected(self, client):
+        with pytest.raises(ValueError, match="Invalid prompt_id"):
+            await client.delete_queue_item("../../free")
+
+    async def test_path_segment_traversal_rejected(self, client):
+        with pytest.raises(ValueError, match="invalid characters"):
+            await client.get_object_info("../../userdata")
+
+    async def test_path_segment_empty_rejected(self, client):
+        with pytest.raises(ValueError, match="must not be empty"):
+            await client.delete_download_task("")
+
+    async def test_http_method_injection_rejected(self, client):
+        with pytest.raises(ValueError, match="HTTP method not allowed"):
+            await client._request("CONNECT", "/queue")
+
+
 class TestModelManagerClient:
     @respx.mock
     async def test_get_model_manager_folders(self, client):

tests/test_integration.py

@@ -1,22 +1,28 @@
 """End-to-end integration tests with mocked ComfyUI backend.
 
-These tests wire up the full server stack (config -> security -> tools -> client)
+These tests wire up tool registration directly (via register_*_tools return dicts)
 and exercise tools through the same code paths used in production.
 """
 
 import json
 
 import httpx
 import respx
+from mcp.server.fastmcp import FastMCP
 
-from comfyui_mcp.config import ComfyUISettings, SecuritySettings, Settings
-from comfyui_mcp.security.inspector import WorkflowBlockedError
-from comfyui_mcp.server import _build_server
+from comfyui_mcp.audit import AuditLogger
+from comfyui_mcp.client import ComfyUIClient
+from comfyui_mcp.security.inspector import WorkflowBlockedError, WorkflowInspector
+from comfyui_mcp.security.rate_limit import RateLimiter
+from comfyui_mcp.security.sanitizer import PathSanitizer
+from comfyui_mcp.tools.discovery import register_discovery_tools
+from comfyui_mcp.tools.generation import register_generation_tools
+from comfyui_mcp.tools.jobs import register_job_tools
 
 
 class TestImageGenerationFlow:
     @respx.mock
-    async def test_generate_image_lists_models_then_generates(self):
+    async def test_generate_image_lists_models_then_generates(self, tmp_path):
         """Full flow: list models -> generate image -> check job."""
         respx.get("http://mock-comfyui:8188/models/checkpoints").mock(
             return_value=httpx.Response(200, json=["sd_v15.safetensors"])
@@ -47,58 +53,66 @@ async def test_generate_image_lists_models_then_generates(self):
             )
         )
 
-        settings = Settings(comfyui=ComfyUISettings(url="http://mock-comfyui:8188"))
-        server, _ = _build_server(settings)
+        client = ComfyUIClient(base_url="http://mock-comfyui:8188")
+        audit = AuditLogger(audit_file=tmp_path / "audit.log")
+        limiter = RateLimiter(max_per_minute=60)
+        inspector = WorkflowInspector(mode="audit", dangerous_nodes=[], allowed_nodes=[])
+        sanitizer = PathSanitizer(allowed_extensions=[".png", ".jpg", ".jpeg", ".webp"])
+        mcp = FastMCP("test")
+
+        discovery_tools = register_discovery_tools(mcp, client, audit, limiter, sanitizer)
+        gen_tools = register_generation_tools(mcp, client, audit, limiter, inspector)
+        job_tools = register_job_tools(mcp, client, audit, limiter)
 
         # Step 1: Discover available models
-        tools = server._tool_manager._tools
-        list_models_fn = tools["list_models"].fn
-        models = await list_models_fn(folder="checkpoints")
+        models = await discovery_tools["list_models"](folder="checkpoints")
         assert "sd_v15.safetensors" in models
 
         # Step 2: Generate an image
-        generate_fn = tools["generate_image"].fn
-        result = await generate_fn(prompt="a sunset over mountains")
+        result = await gen_tools["generate_image"](prompt="a sunset over mountains")
         assert "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" in result
 
         # Step 3: Check the job
-        get_job_fn = tools["get_job"].fn
-        job = await get_job_fn(prompt_id="aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee")
+        job = await job_tools["get_job"](prompt_id="aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee")
         assert "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" in job
 
     @respx.mock
-    async def test_run_workflow_with_dangerous_node_in_audit_mode(self):
+    async def test_run_workflow_with_dangerous_node_in_audit_mode(self, tmp_path):
         """Audit mode logs dangerous nodes but still submits the workflow."""
         respx.post("http://mock-comfyui:8188/prompt").mock(
             return_value=httpx.Response(
                 200, json={"prompt_id": "11111111-2222-3333-4444-555555555555"}
             )
         )
 
-        settings = Settings(comfyui=ComfyUISettings(url="http://mock-comfyui:8188"))
-        server, _ = _build_server(settings)
+        client = ComfyUIClient(base_url="http://mock-comfyui:8188")
+        audit = AuditLogger(audit_file=tmp_path / "audit.log")
+        limiter = RateLimiter(max_per_minute=60)
+        inspector = WorkflowInspector(mode="audit", dangerous_nodes=["Terminal"], allowed_nodes=[])
+        mcp = FastMCP("test")
 
-        run_workflow_fn = server._tool_manager._tools["run_workflow"].fn
+        tools = register_generation_tools(mcp, client, audit, limiter, inspector)
         workflow = json.dumps({"1": {"class_type": "Terminal", "inputs": {}}})
-        result = await run_workflow_fn(workflow=workflow)
+        result = await tools["run_workflow"](workflow=workflow)
         assert "11111111-2222-3333-4444-555555555555" in result
         assert "Terminal" in result
 
-    async def test_run_workflow_blocked_in_enforce_mode(self):
+    async def test_run_workflow_blocked_in_enforce_mode(self, tmp_path):
         """Enforce mode blocks workflows with unapproved nodes."""
-        settings = Settings(
-            comfyui=ComfyUISettings(url="http://mock-comfyui:8188"),
-            security=SecuritySettings(
-                mode="enforce",
-                allowed_nodes=["KSampler", "CLIPTextEncode"],
-            ),
+        client = ComfyUIClient(base_url="http://mock-comfyui:8188")
+        audit = AuditLogger(audit_file=tmp_path / "audit.log")
+        limiter = RateLimiter(max_per_minute=60)
+        inspector = WorkflowInspector(
+            mode="enforce",
+            dangerous_nodes=[],
+            allowed_nodes=["KSampler", "CLIPTextEncode"],
         )
-        server, _ = _build_server(settings)
+        mcp = FastMCP("test")
 
-        run_workflow_fn = server._tool_manager._tools["run_workflow"].fn
+        tools = register_generation_tools(mcp, client, audit, limiter, inspector)
         workflow = json.dumps({"1": {"class_type": "MaliciousNode", "inputs": {}}})
         try:
-            await run_workflow_fn(workflow=workflow)
+            await tools["run_workflow"](workflow=workflow)
             raise AssertionError("Should have raised WorkflowBlockedError")
         except WorkflowBlockedError as e:
             assert "MaliciousNode" in str(e)

tests/test_model_manager.py

@@ -1,9 +1,15 @@
+import time
+
 import httpx
 import pytest
 import respx
 
 from comfyui_mcp.client import ComfyUIClient
-from comfyui_mcp.model_manager import ModelManagerDetector, ModelManagerUnavailableError
+from comfyui_mcp.model_manager import (
+    _NEGATIVE_TTL_SECONDS,
+    ModelManagerDetector,
+    ModelManagerUnavailableError,
+)
 
 
 @pytest.fixture
@@ -103,3 +109,31 @@ async def test_validate_folder_unknown(self, client):
         detector = ModelManagerDetector(client)
         with pytest.raises(ValueError, match="not a valid model folder"):
             await detector.validate_folder("invalid_folder")
+
+    @respx.mock
+    async def test_negative_cache_expires(self, client):
+        """After a failed probe, re-probes once the negative TTL expires."""
+        route = respx.get("http://test:8188/model-manager/models")
+
+        # First call: Model Manager not available
+        route.mock(return_value=httpx.Response(404))
+        detector = ModelManagerDetector(client)
+        assert await detector.is_available() is False
+        assert route.call_count == 1
+
+        # Second call within TTL: should use cached negative result
+        route.mock(
+            return_value=httpx.Response(
+                200,
+                json={"success": True, "data": {"checkpoints": ["/models/checkpoints"]}},
+            )
+        )
+        assert await detector.is_available() is False
+        assert route.call_count == 1  # No new probe
+
+        # Simulate TTL expiry
+        detector._last_failure = time.monotonic() - _NEGATIVE_TTL_SECONDS - 1
+
+        # Third call after TTL: should re-probe and succeed
+        assert await detector.is_available() is True
+        assert route.call_count == 2