From 5a49b9782159c01c65e0ca6d5b2c80f9de779a3b Mon Sep 17 00:00:00 2001
From: Adam Outler <adamoutler@gmail.com>
Date: Sat, 31 Jan 2026 02:54:00 +0000
Subject: [PATCH] Fixes for Coderabbit review

---
 server/api_server/api_server_start.py         |  3 +
 server/api_server/openapi/schemas.py          |  1 +
 server/api_server/openapi/spec_generator.py   |  6 +-
 test/api_endpoints/test_mcp_disabled_tools.py | 63 +++++++++++++++++++
 4 files changed, 70 insertions(+), 3 deletions(-)
 create mode 100644 test/api_endpoints/test_mcp_disabled_tools.py

diff --git a/server/api_server/api_server_start.py b/server/api_server/api_server_start.py
index 31f6916c..de68c90d 100755
--- a/server/api_server/api_server_start.py
+++ b/server/api_server/api_server_start.py
@@ -7,6 +7,7 @@ import os
 from flask import Flask, request, jsonify, Response
 from models.device_instance import DeviceInstance  # noqa: E402
 from flask_cors import CORS
+from werkzeug.exceptions import HTTPException
 
 # Register NetAlertX directories
 INSTALL_PATH = os.getenv("NETALERTX_APP", "/app")
@@ -105,6 +106,8 @@ app = Flask(__name__)
 @app.errorhandler(Exception)
 def handle_500_error(e):
     """Global error handler for uncaught exceptions."""
+    if isinstance(e, HTTPException):
+        return e
     mylog("none", [f"[API] Uncaught exception: {e}"])
     return jsonify({
         "success": False,
diff --git a/server/api_server/openapi/schemas.py b/server/api_server/openapi/schemas.py
index 3e7e6add..118763d7 100644
--- a/server/api_server/openapi/schemas.py
+++ b/server/api_server/openapi/schemas.py
@@ -679,6 +679,7 @@ class CreateEventRequest(BaseModel):
 # =============================================================================
 # SESSIONS SCHEMAS
 # =============================================================================
+class SessionInfo(BaseModel):
     """Session information."""
     model_config = ConfigDict(
         extra="allow",
diff --git a/server/api_server/openapi/spec_generator.py b/server/api_server/openapi/spec_generator.py
index 94b7099e..3d42779f 100644
--- a/server/api_server/openapi/spec_generator.py
+++ b/server/api_server/openapi/spec_generator.py
@@ -77,7 +77,7 @@ def generate_openapi_spec(
             # Apply default disabled tools from setting `MCP_DISABLED_TOOLS`, env var, or hard-coded defaults
             # Format: comma-separated operation IDs, e.g. "dbquery_read,dbquery_write"
             try:
-                disabled_env = ""
+                disabled_env = None
                 # Prefer setting from app.conf/settings when available
                 try:
                     from helper import get_setting_value
@@ -88,9 +88,9 @@ def generate_openapi_spec(
                     # If helper is unavailable, fall back to environment
                     pass
 
-                if disabled_env is None:
+                if not disabled_env:
                     env_val = os.getenv("MCP_DISABLED_TOOLS")
-                    if env_val is not None:
+                    if env_val:
                         disabled_env = env_val.strip()
 
                 # If still not set, apply safe hard-coded defaults
diff --git a/test/api_endpoints/test_mcp_disabled_tools.py b/test/api_endpoints/test_mcp_disabled_tools.py
new file mode 100644
index 00000000..fd9e67b2
--- /dev/null
+++ b/test/api_endpoints/test_mcp_disabled_tools.py
@@ -0,0 +1,63 @@
+
+import os
+import sys
+import pytest
+from unittest.mock import patch, MagicMock
+
+# Use cwd as fallback if env var is not set, assuming running from project root
+INSTALL_PATH = os.getenv('NETALERTX_APP', os.getcwd())
+sys.path.extend([f"{INSTALL_PATH}/front/plugins", f"{INSTALL_PATH}/server"])
+
+from api_server.openapi.spec_generator import generate_openapi_spec
+from api_server.api_server_start import app
+
+class TestMCPDisabledTools:
+    
+    def test_disabled_tools_via_env_var(self):
+        """Test that MCP_DISABLED_TOOLS env var disables specific tools."""
+        # Clean registry first to ensure clean state
+        from api_server.openapi.registry import clear_registry
+        clear_registry()
+
+        # Mock get_setting_value to return None (simulating no config setting)
+        # and mock os.getenv to return our target list
+        with patch("helper.get_setting_value", return_value=None), \
+             patch.dict(os.environ, {"MCP_DISABLED_TOOLS": "search_devices_api"}):
+            
+            spec = generate_openapi_spec(flask_app=app)
+            
+            # Locate the operation
+            # search_devices_api is usually mapped to /devices/search [POST] or similar
+            # We search the spec for the operationId
+            
+            found = False
+            for path, methods in spec["paths"].items():
+                for method, op in methods.items():
+                    if op["operationId"] == "search_devices_api":
+                        assert op.get("x-mcp-disabled") is True
+                        found = True
+            
+            assert found, "search_devices_api operation not found in spec"
+
+    def test_disabled_tools_default_fallback(self):
+        """Test fallback to defaults when no setting or env var exists."""
+        from api_server.openapi.registry import clear_registry
+        clear_registry()
+
+        with patch("helper.get_setting_value", return_value=None), \
+             patch.dict(os.environ, {}, clear=True): # Clear env to ensure no MCP_DISABLED_TOOLS
+            
+            spec = generate_openapi_spec(flask_app=app)
+            
+            # Default is "dbquery_read,dbquery_write"
+            
+            # Check dbquery_read
+            found_read = False
+            for path, methods in spec["paths"].items():
+                for method, op in methods.items():
+                    if op["operationId"] == "dbquery_read":
+                        assert op.get("x-mcp-disabled") is True
+                        found_read = True
+            
+            assert found_read, "dbquery_read should be disabled by default"
+