"""Simple API key authentication middleware for MCP HTTP."""
from contextvars import ContextVar
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.status import HTTP_401_UNAUTHORIZED
# Context variables to store LangSmith config for current request
# These are used to pass config from middleware to FastMCP's session state
api_key_context: ContextVar[str] = ContextVar("api_key", default="")
workspace_id_context: ContextVar[str] = ContextVar("workspace_id", default="")
endpoint_context: ContextVar[str] = ContextVar("endpoint", default="")
def get_api_key() -> str:
"""Get the API key from the current request context."""
return api_key_context.get("")
class APIKeyMiddleware(BaseHTTPMiddleware):
"""
Middleware that extracts LANGSMITH-API-KEY header and optional config headers.
FastMCP will handle session management automatically. We extract:
1. LANGSMITH-API-KEY (required)
2. LANGSMITH-WORKSPACE-ID (optional)
3. LANGSMITH-ENDPOINT (optional)
""" # noqa: W293
async def dispatch(self, request: Request, call_next):
# Skip authentication for health check
if request.url.path == "/health":
return await call_next(request)
# Require LANGSMITH-API-KEY header
api_key = request.headers.get("LANGSMITH-API-KEY")
if not api_key:
return JSONResponse(
status_code=HTTP_401_UNAUTHORIZED,
content={"error": "Missing LANGSMITH-API-KEY header"},
)
# Get optional headers
workspace_id = request.headers.get("LANGSMITH-WORKSPACE-ID", "")
endpoint = request.headers.get("LANGSMITH-ENDPOINT", "")
# Store in request state and context variables
# FastMCP will handle storing this in session state via the Context
request.state.api_key = api_key
request.state.workspace_id = workspace_id
request.state.endpoint = endpoint
api_key_context.set(api_key)
workspace_id_context.set(workspace_id)
endpoint_context.set(endpoint)
try:
return await call_next(request)
finally:
# Clear context after request
api_key_context.set("")
workspace_id_context.set("")
endpoint_context.set("")
