feat(ci): add MCP protocol conformance tests

.github/workflows/conformance.yaml

@@ -0,0 +1,77 @@
+name: MCP Conformance
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+permissions:
+  contents: read
+
+jobs:
+  conformance:
+    runs-on: ubuntu-latest
+    name: Protocol conformance (streamable HTTP)
+
+    steps:
+      - uses: actions/checkout@v7
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+
+      - name: Install package
+        run: uv sync --all-extras
+
+      # Start the server over streamable HTTP with no kubeconfig / no cluster.
+      # health_check and all tools degrade gracefully when K8s is unreachable,
+      # so the server starts and answers initialize/tools-list without a cluster.
+      - name: Start MCP server
+        env:
+          KUBECONFIG: /nonexistent
+        run: |
+          uv run kubeflow-mcp serve --transport http --no-banner > server.log 2>&1 &
+          echo $! > server.pid
+
+      # FastMCP streamable-HTTP has no plain health endpoint, so probe the MCP
+      # endpoint with a real initialize request and wait for HTTP 200.
+      - name: Wait for server
+        run: |
+          timeout 60 bash -c 'until curl -sf -o /dev/null \
+            http://localhost:8000/mcp \
+            -X POST \
+            -H "Content-Type: application/json" \
+            -H "Accept: application/json, text/event-stream" \
+            -d "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":\"2025-06-18\",\"capabilities\":{},\"clientInfo\":{\"name\":\"ci\",\"version\":\"1\"}}}"; \
+            do sleep 1; done'
+
+      - name: Run MCP conformance suite
+        uses: modelcontextprotocol/conformance@v0.1.16
+        with:
+          mode: server
+          url: http://localhost:8000/mcp
+          suite: active
+          expected-failures: tests/conformance/expected-failures.yaml
+          node-version: 22
+
+      - name: Dump server logs
+        if: always()
+        run: cat server.log || true
+
+      - name: Stop MCP server
+        if: always()
+        run: |
+          if [ -f server.pid ]; then
+            kill "$(cat server.pid)" 2>/dev/null || true
+          fi

CONTRIBUTING.md

@@ -59,6 +59,41 @@ To run unit tests locally, use the following `make` command:
 make test-python
 ```
 
+### Conformance Testing
+
+We validate that the server correctly implements the MCP protocol
+(`initialize`, `tools/list`, `tools/call`, etc.) using the official
+[MCP conformance framework](https://github.com/modelcontextprotocol/conformance).
+This runs in CI on every PR (`.github/workflows/conformance.yaml`) and needs
+**no Kubernetes cluster and no LLM** — the server starts against an unreachable
+cluster and its tools degrade gracefully.
+
+To run the same suite locally (requires Node.js 20+):
+
+```bash
+make conformance
+```
+
+This starts `kubeflow-mcp serve --transport http` in the background, waits for
+it to answer `initialize`, then runs the conformance `active` suite against
+`http://localhost:8000/mcp`.
+
+**Baseline file.** `tests/conformance/expected-failures.yaml` lists scenarios
+that are allowed to fail — these are known spec gaps for this server (mostly
+scenarios that assume the reference server's test fixtures, plus a few
+unadvertised capabilities). Every scenario *not* listed must pass, and CI also
+fails if a listed scenario unexpectedly starts passing. When you add protocol
+features, remove the now-passing scenarios from the baseline; if you
+intentionally change behavior, update the list with a one-line reason.
+
+**DNS rebinding protection.** The HTTP/SSE transports validate the `Host` and
+`Origin` headers (`kubeflow_mcp/core/transport_security.py`) so a locally-bound
+server cannot be reached by a malicious web page via DNS rebinding. Protection
+is on by default and allows loopback hosts; override the allowlists with
+`KUBEFLOW_MCP_ALLOWED_HOSTS` / `KUBEFLOW_MCP_ALLOWED_ORIGINS` (comma-separated,
+`:*` port wildcard supported) when serving on a non-loopback address, or set
+`KUBEFLOW_MCP_DNS_REBINDING_PROTECTION=false` to disable it (not recommended).
+
 ## Commit Messages
 
 We use [Conventional Commits](https://www.conventionalcommits.org/):

Makefile

@@ -62,6 +62,18 @@ test-cov: ## Run tests with HTML coverage report
 
 ##@ Dev Tools
 
+.PHONY: conformance
+conformance: install-dev ## Run MCP protocol conformance suite against a local HTTP server
+	@echo "Starting kubeflow-mcp on http://localhost:8000/mcp (no cluster required)..."
+	@KUBECONFIG=/nonexistent uv run kubeflow-mcp serve --transport http --no-banner > /tmp/kubeflow-mcp-conformance.log 2>&1 & \
+	  SERVER_PID=$$!; \
+	  trap "kill $$SERVER_PID 2>/dev/null || true" EXIT; \
+	  timeout 60 bash -c 'until curl -sf -o /dev/null http://localhost:8000/mcp -X POST -H "Content-Type: application/json" -H "Accept: application/json, text/event-stream" -d "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":\"2025-06-18\",\"capabilities\":{},\"clientInfo\":{\"name\":\"make\",\"version\":\"1\"}}}"; do sleep 1; done'; \
+	  npx -y @modelcontextprotocol/conformance@0.1.16 server \
+	    --url http://localhost:8000/mcp \
+	    --suite active \
+	    --expected-failures tests/conformance/expected-failures.yaml
+
 TRANSPORT ?= stdio
 
 inspector: install-dev ## Launch MCP Inspector (TRANSPORT=stdio|http|sse)

README.md

@@ -63,6 +63,9 @@ The server listens on `http://localhost:8000/mcp`.
 | `KUBEFLOW_MCP_JWT_AUDIENCE` | _(none)_ | Expected JWT audience |
 | `KUBEFLOW_MCP_CLIENTS` | `trainer` | Comma-separated client modules to load |
 | `KUBEFLOW_MCP_PERSONA` | `readonly` | Tool persona (`readonly`, `data-scientist`, `ml-engineer`, `platform-admin`) |
+| `KUBEFLOW_MCP_ALLOWED_HOSTS` | _(loopback)_ | Comma-separated `Host` header allowlist for DNS rebinding protection; `:*` port wildcard supported (e.g. `mcp.example.com,mcp.example.com:*`) |
+| `KUBEFLOW_MCP_ALLOWED_ORIGINS` | _(loopback)_ | Comma-separated `Origin` header allowlist; `:*` port wildcard supported (e.g. `https://mcp.example.com`) |
+| `KUBEFLOW_MCP_DNS_REBINDING_PROTECTION` | `true` | Set `false` to disable Host/Origin validation (not recommended) |
 | `LOG_FORMAT` | `json` | Log format (`json`, `console`) |
 | `LOG_LEVEL` | `INFO` | Log level (`DEBUG`, `INFO`, `WARNING`, `ERROR`) |
 
@@ -81,6 +84,8 @@ The server listens on `http://localhost:8000/mcp`.
 
 For in-cluster deployments, replace `localhost:8000` with the Kubernetes Service address and mount `KUBEFLOW_MCP_AUTH_TOKEN` from a Secret.
 
+> **Note:** DNS rebinding protection allows only loopback `Host`/`Origin` headers by default. When exposing the server through a Service or Ingress, set `KUBEFLOW_MCP_ALLOWED_HOSTS` (e.g. `KUBEFLOW_MCP_ALLOWED_HOSTS=kubeflow-mcp.kubeflow.svc:*,mcp.example.com`) or requests will be rejected with HTTP 421.
+
 ### Example: Fine-tune a model via AI agent
 
 Once connected, your AI agent can run a complete training workflow through natural language:

kubeflow_mcp/cli.py

@@ -163,6 +163,7 @@ def serve(
     )
 
     auth_provider = None
+    http_middleware = None
     if transport != "stdio":
         auth_provider = build_auth_provider(cfg.auth)
         if auth_provider is None:
@@ -172,6 +173,27 @@ def serve(
                 "or KUBEFLOW_MCP_JWKS_URI for JWT verification."
             )
 
+        from starlette.middleware import Middleware
+
+        from kubeflow_mcp.core.transport_security import (
+            DNSRebindingProtectionMiddleware,
+            build_transport_security_settings,
+        )
+
+        security_settings = build_transport_security_settings(cfg.security)
+        if security_settings.enable_dns_rebinding_protection:
+            http_middleware = [
+                Middleware(
+                    DNSRebindingProtectionMiddleware,
+                    settings=security_settings,
+                )
+            ]
+        else:
+            logger.warning(
+                "DNS rebinding protection is disabled — set "
+                "KUBEFLOW_MCP_DNS_REBINDING_PROTECTION=true to re-enable it."
+            )
+
     client_list = [c.strip() for c in clients.split(",")]
     server = create_server(
         clients=client_list,
@@ -185,9 +207,13 @@ def serve(
     if transport == "stdio":
         server.run(show_banner=show_banner)
     elif transport == "sse":
-        server.run(transport="sse", show_banner=show_banner)
+        server.run(transport="sse", show_banner=show_banner, middleware=http_middleware)
     else:
-        server.run(transport="streamable-http", show_banner=show_banner)
+        server.run(
+            transport="streamable-http",
+            show_banner=show_banner,
+            middleware=http_middleware,
+        )
 
 
 @cli.command()

kubeflow_mcp/core/config.py

@@ -115,6 +115,25 @@ class AuthConfig(BaseModel):
     audience: str | None = Field(default=None)
 
 
+class SecurityConfig(BaseModel):
+    """HTTP transport security (DNS rebinding protection).
+
+    Validates the ``Host`` and ``Origin`` headers of incoming HTTP requests so a
+    malicious web page cannot use DNS rebinding to reach a locally-bound server.
+    Only applies to the ``http``/``sse`` transports; ``stdio`` has no HTTP surface.
+
+    ``allowed_hosts``/``allowed_origins`` support an exact match or a ``:*`` port
+    wildcard (e.g. ``localhost:*``). When left empty they default to loopback
+    values (IPv4 ``127.0.0.1``, IPv6 ``[::1]``, and ``localhost``), which is
+    correct for the default ``127.0.0.1`` bind used in dev and CI. IPv6 hosts
+    must be given in bracketed form (e.g. ``[::1]:*``) to match the Host header.
+    """
+
+    enable_dns_rebinding_protection: bool = Field(default=True)
+    allowed_hosts: list[str] = Field(default_factory=list)
+    allowed_origins: list[str] = Field(default_factory=list)
+
+
 class ResilienceConfig(BaseModel):
     """Rate limiter and circuit breaker tuning."""
 
@@ -149,6 +168,7 @@ class Config(BaseModel):
 
     server: ServerConfig = Field(default_factory=ServerConfig)
     auth: AuthConfig = Field(default_factory=AuthConfig)
+    security: SecurityConfig = Field(default_factory=SecurityConfig)
     resilience: ResilienceConfig = Field(default_factory=ResilienceConfig)
     trainer: TrainerConfig = Field(default_factory=TrainerConfig)
     optimizer: OptimizerConfig = Field(default_factory=OptimizerConfig)
@@ -272,6 +292,29 @@ def load_config(config_path: Path | None = None) -> Config:
         audience=os.getenv("KUBEFLOW_MCP_JWT_AUDIENCE", auth_file.get("audience")),
     )
 
+    security_file = file_config.get("security", {})
+
+    def _csv_env(name: str, file_default: list[str]) -> list[str]:
+        raw = os.getenv(name)
+        if raw is None:
+            return list(file_default)
+        return [item.strip() for item in raw.split(",") if item.strip()]
+
+    dns_protect_env = os.getenv("KUBEFLOW_MCP_DNS_REBINDING_PROTECTION")
+    security = SecurityConfig(
+        enable_dns_rebinding_protection=(
+            dns_protect_env.strip().lower() not in {"0", "false", "no", "off"}
+            if dns_protect_env is not None
+            else security_file.get("enable_dns_rebinding_protection", True)
+        ),
+        allowed_hosts=_csv_env(
+            "KUBEFLOW_MCP_ALLOWED_HOSTS", security_file.get("allowed_hosts", [])
+        ),
+        allowed_origins=_csv_env(
+            "KUBEFLOW_MCP_ALLOWED_ORIGINS", security_file.get("allowed_origins", [])
+        ),
+    )
+
     resilience_file = file_config.get("resilience", {})
     resilience = ResilienceConfig(
         rate_limit=float(
@@ -297,6 +340,7 @@ def load_config(config_path: Path | None = None) -> Config:
     return Config(
         server=server,
         auth=auth,
+        security=security,
         resilience=resilience,
         trainer=trainer,
         optimizer=optimizer,

kubeflow_mcp/core/middleware.py

@@ -69,7 +69,7 @@ class AuditIdentityMiddleware:
 
         from kubeflow_mcp.core.middleware import AuditIdentityMiddleware
         mcp = FastMCP("kubeflow-mcp-server")
-        mcp.add_middleware(AuditIdentityMiddleware)
+        mcp.add_middleware(AuditIdentityMiddleware())
     """
 
     async def __call__(self, context: Any, call_next: Any) -> Any:

kubeflow_mcp/core/server.py

@@ -362,7 +362,7 @@ def create_server(  # noqa: C901
     # Bridge FastMCP async context into sync _audit_wrap via ContextVars
     from kubeflow_mcp.core.middleware import AuditIdentityMiddleware
 
-    mcp.add_middleware(AuditIdentityMiddleware)
+    mcp.add_middleware(AuditIdentityMiddleware())
 
     # Merge tool metadata from client modules
     all_descriptions: dict[str, str] = {}