Included new oauth, middleware, and beta features.

Author: Knucklessg1

.pre-commit-config.yaml

@@ -27,7 +27,7 @@ repos:
     hooks:
     - id: black
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.13.3
+    rev: v0.14.1
     hooks:
       - id: ruff
         types_or: [ python, pyi, jupyter ]

requirements.txt

@@ -1,4 +1,4 @@
-fastmcp>=2.11.3
+fastmcp>=2.12.4
 sentence_transformers>=5.1.0
 markdownify>=1.2.0
 beautifulsoup4>=4.13.5

vector_mcp/vector_mcp.py

@@ -6,8 +6,15 @@
 import logging
 from pathlib import Path
 from typing import Optional, Dict, Union
-from fastmcp import FastMCP, Context
 from pydantic import Field
+from fastmcp import FastMCP, Context
+from fastmcp.server.auth.oidc_proxy import OIDCProxy
+from fastmcp.server.auth import OAuthProxy, RemoteAuthProvider
+from fastmcp.server.auth.providers.jwt import JWTVerifier, StaticTokenVerifier
+from fastmcp.server.middleware.logging import LoggingMiddleware
+from fastmcp.server.middleware.timing import TimingMiddleware
+from fastmcp.server.middleware.rate_limiting import RateLimitingMiddleware
+from fastmcp.server.middleware.error_handling import ErrorHandlingMiddleware
 from vector_mcp.retriever.retriever import RAGRetriever
 from vector_mcp.retriever.pgvector_retriever import PGVectorRetriever
 from vector_mcp.retriever.qdrant_retriever import QdrantRetriever
@@ -558,20 +565,225 @@ def vector_mcp():
         default=8000,
         help="Port number for HTTP transport (default: 8000)",
     )
+    parser.add_argument(
+        "--auth-type",
+        default="none",
+        choices=["none", "static", "jwt", "oauth-proxy", "oidc-proxy", "remote-oauth"],
+        help="Authentication type for MCP server: 'none' (disabled), 'static' (internal), 'jwt' (external token verification), 'oauth-proxy', 'oidc-proxy', 'remote-oauth' (external) (default: none)",
+    )
+    # JWT/Token params
+    parser.add_argument(
+        "--token-jwks-uri", default=None, help="JWKS URI for JWT verification"
+    )
+    parser.add_argument(
+        "--token-issuer", default=None, help="Issuer for JWT verification"
+    )
+    parser.add_argument(
+        "--token-audience", default=None, help="Audience for JWT verification"
+    )
+    # OAuth Proxy params
+    parser.add_argument(
+        "--oauth-upstream-auth-endpoint",
+        default=None,
+        help="Upstream authorization endpoint for OAuth Proxy",
+    )
+    parser.add_argument(
+        "--oauth-upstream-token-endpoint",
+        default=None,
+        help="Upstream token endpoint for OAuth Proxy",
+    )
+    parser.add_argument(
+        "--oauth-upstream-client-id",
+        default=None,
+        help="Upstream client ID for OAuth Proxy",
+    )
+    parser.add_argument(
+        "--oauth-upstream-client-secret",
+        default=None,
+        help="Upstream client secret for OAuth Proxy",
+    )
+    parser.add_argument(
+        "--oauth-base-url", default=None, help="Base URL for OAuth Proxy"
+    )
+    # OIDC Proxy params
+    parser.add_argument(
+        "--oidc-config-url", default=None, help="OIDC configuration URL"
+    )
+    parser.add_argument("--oidc-client-id", default=None, help="OIDC client ID")
+    parser.add_argument("--oidc-client-secret", default=None, help="OIDC client secret")
+    parser.add_argument("--oidc-base-url", default=None, help="Base URL for OIDC Proxy")
+    # Remote OAuth params
+    parser.add_argument(
+        "--remote-auth-servers",
+        default=None,
+        help="Comma-separated list of authorization servers for Remote OAuth",
+    )
+    parser.add_argument(
+        "--remote-base-url", default=None, help="Base URL for Remote OAuth"
+    )
+    # Common
+    parser.add_argument(
+        "--allowed-client-redirect-uris",
+        default=None,
+        help="Comma-separated list of allowed client redirect URIs",
+    )
+    # Eunomia params
+    parser.add_argument(
+        "--eunomia-type",
+        default="none",
+        choices=["none", "embedded", "remote"],
+        help="Eunomia authorization type: 'none' (disabled), 'embedded' (built-in), 'remote' (external) (default: none)",
+    )
+    parser.add_argument(
+        "--eunomia-policy-file",
+        default="mcp_policies.json",
+        help="Policy file for embedded Eunomia (default: mcp_policies.json)",
+    )
+    parser.add_argument(
+        "--eunomia-remote-url", default=None, help="URL for remote Eunomia server"
+    )
 
     args = parser.parse_args()
 
     if args.port < 0 or args.port > 65535:
         print(f"Error: Port {args.port} is out of valid range (0-65535).")
         sys.exit(1)
 
+    # Set auth based on type
+    auth = None
+    allowed_uris = (
+        args.allowed_client_redirect_uris.split(",")
+        if args.allowed_client_redirect_uris
+        else None
+    )
+
+    if args.auth_type == "none":
+        auth = None
+    elif args.auth_type == "static":
+        # Internal static tokens (hardcoded example)
+        auth = StaticTokenVerifier(
+            tokens={
+                "test-token": {"client_id": "test-user", "scopes": ["read", "write"]},
+                "admin-token": {"client_id": "admin", "scopes": ["admin"]},
+            }
+        )
+    elif args.auth_type == "jwt":
+        if not (args.token_jwks_uri and args.token_issuer and args.token_audience):
+            print(
+                "Error: jwt requires --token-jwks-uri, --token-issuer, --token-audience"
+            )
+            sys.exit(1)
+        auth = JWTVerifier(
+            jwks_uri=args.token_jwks_uri,
+            issuer=args.token_issuer,
+            audience=args.token_audience,
+        )
+    elif args.auth_type == "oauth-proxy":
+        if not (
+            args.oauth_upstream_auth_endpoint
+            and args.oauth_upstream_token_endpoint
+            and args.oauth_upstream_client_id
+            and args.oauth_upstream_client_secret
+            and args.oauth_base_url
+            and args.token_jwks_uri
+            and args.token_issuer
+            and args.token_audience
+        ):
+            print(
+                "Error: oauth-proxy requires --oauth-upstream-auth-endpoint, --oauth-upstream-token-endpoint, --oauth-upstream-client-id, --oauth-upstream-client-secret, --oauth-base-url, --token-jwks-uri, --token-issuer, --token-audience"
+            )
+            sys.exit(1)
+        token_verifier = JWTVerifier(
+            jwks_uri=args.token_jwks_uri,
+            issuer=args.token_issuer,
+            audience=args.token_audience,
+        )
+        auth = OAuthProxy(
+            upstream_authorization_endpoint=args.oauth_upstream_auth_endpoint,
+            upstream_token_endpoint=args.oauth_upstream_token_endpoint,
+            upstream_client_id=args.oauth_upstream_client_id,
+            upstream_client_secret=args.oauth_upstream_client_secret,
+            token_verifier=token_verifier,
+            base_url=args.oauth_base_url,
+            allowed_client_redirect_uris=allowed_uris,
+        )
+    elif args.auth_type == "oidc-proxy":
+        if not (
+            args.oidc_config_url
+            and args.oidc_client_id
+            and args.oidc_client_secret
+            and args.oidc_base_url
+        ):
+            print(
+                "Error: oidc-proxy requires --oidc-config-url, --oidc-client-id, --oidc-client-secret, --oidc-base-url"
+            )
+            sys.exit(1)
+        auth = OIDCProxy(
+            config_url=args.oidc_config_url,
+            client_id=args.oidc_client_id,
+            client_secret=args.oidc_client_secret,
+            base_url=args.oidc_base_url,
+            allowed_client_redirect_uris=allowed_uris,
+        )
+    elif args.auth_type == "remote-oauth":
+        if not (
+            args.remote_auth_servers
+            and args.remote_base_url
+            and args.token_jwks_uri
+            and args.token_issuer
+            and args.token_audience
+        ):
+            print(
+                "Error: remote-oauth requires --remote-auth-servers, --remote-base-url, --token-jwks-uri, --token-issuer, --token-audience"
+            )
+            sys.exit(1)
+        auth_servers = [url.strip() for url in args.remote_auth_servers.split(",")]
+        token_verifier = JWTVerifier(
+            jwks_uri=args.token_jwks_uri,
+            issuer=args.token_issuer,
+            audience=args.token_audience,
+        )
+        auth = RemoteAuthProvider(
+            token_verifier=token_verifier,
+            authorization_servers=auth_servers,
+            base_url=args.remote_base_url,
+        )
+    mcp.auth = auth
+    if args.eunomia_type != "none":
+        from eunomia_mcp import create_eunomia_middleware
+
+        if args.eunomia_type == "embedded":
+            if not args.eunomia_policy_file:
+                print("Error: embedded Eunomia requires --eunomia-policy-file")
+                sys.exit(1)
+            middleware = create_eunomia_middleware(policy_file=args.eunomia_policy_file)
+            mcp.add_middleware(middleware)
+        elif args.eunomia_type == "remote":
+            if not args.eunomia_remote_url:
+                print("Error: remote Eunomia requires --eunomia-remote-url")
+                sys.exit(1)
+            middleware = create_eunomia_middleware(
+                use_remote_eunomia=args.eunomia_remote_url
+            )
+            mcp.add_middleware(middleware)
+
+    mcp.add_middleware(
+        ErrorHandlingMiddleware(include_traceback=True, transform_errors=True)
+    )
+    mcp.add_middleware(
+        RateLimitingMiddleware(max_requests_per_second=10.0, burst_capacity=20)
+    )
+    mcp.add_middleware(TimingMiddleware())
+    mcp.add_middleware(LoggingMiddleware())
+
     if args.transport == "stdio":
         mcp.run(transport="stdio")
     elif args.transport == "http":
         mcp.run(transport="http", host=args.host, port=args.port)
     elif args.transport == "sse":
         mcp.run(transport="sse", host=args.host, port=args.port)
     else:
+        logger = logging.getLogger("Vector")
         logger.error("Transport not supported")
         sys.exit(1)