Fail fast if Grafana override settings are not provided in sse request

Author: sd2k

src/mcp_grafana/middleware.py

@@ -2,6 +2,7 @@
 
 from mcp.server import FastMCP
 from starlette.datastructures import Headers
+from starlette.exceptions import HTTPException
 
 from .client import GrafanaClient, grafana_client
 from .settings import GrafanaSettings, grafana_settings
@@ -36,8 +37,9 @@ class GrafanaMiddleware:
     This should be used as a context manager before handling the /sse request.
     """
 
-    def __init__(self, request):
+    def __init__(self, request, fail_if_unset=True):
         self.request = request
+        self.fail_if_unset = fail_if_unset
         self.settings_token = None
         self.client_token = None
 
@@ -53,6 +55,8 @@ async def __aenter__(self):
             self.client_token = grafana_client.set(
                 GrafanaClient.from_settings(new_settings)
             )
+        elif self.fail_if_unset:
+            raise HTTPException(status_code=403, detail="No Grafana settings found.")
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         if self.settings_token is not None:

tests/middleware_test.py

@@ -5,6 +5,7 @@
 
 import anyio
 import httpx
+from mcp.server import FastMCP
 from mcp.types import (
     LATEST_PROTOCOL_VERSION,
     CallToolResult,
@@ -20,13 +21,21 @@
 )
 import pytest
 from httpx_sse import aconnect_sse
+from starlette.exceptions import HTTPException
 
-from mcp_grafana import mcp
+from mcp_grafana.tools import add_tools
 from mcp_grafana.middleware import run_sse_async_with_middleware
 
 from pytest_httpserver import HTTPServer
 
 
+@pytest.fixture
+def mcp():
+    mcp = FastMCP("grafana")
+    add_tools(mcp)
+    return mcp
+
+
 class TestMiddleware:
     """
     Test that our injected starlette middleware extracts headers and
@@ -36,7 +45,28 @@ class TestMiddleware:
     """
 
     @pytest.mark.asyncio
-    async def test_multiple_requests(self):
+    async def test_no_headers_provided(self, mcp: FastMCP):
+        """
+        Ensure that the middleware fails if no headers are provided.
+        """
+
+        # Monkeypatch the MCP server to use our middleware.
+        mcp.run_sse_async = MethodType(run_sse_async_with_middleware, mcp)
+        mcp.settings.host = "127.0.0.1"
+        mcp.settings.port = 9500
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(mcp.run_sse_async, name="mcp")
+            # Wait for the server to start.
+            await asyncio.sleep(0.1)
+            client = httpx.AsyncClient(
+                base_url=f"http://{mcp.settings.host}:{mcp.settings.port}"
+            )
+            resp = await client.get("/sse")
+            assert resp.status_code == httpx.codes.FORBIDDEN
+            tg.cancel_scope.cancel()
+
+    @pytest.mark.asyncio
+    async def test_multiple_requests(self, mcp: FastMCP):
         """
         Ensure that the contextvars do not leak across requests.