Spaces:
Sleeping
Sleeping
File size: 2,355 Bytes
1ce695c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | from fastapi import Request, HTTPException
from starlette.middleware.base import BaseHTTPMiddleware
from config.settings import settings
import logging
logger = logging.getLogger(__name__)
class AgentValidatorMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next):
# Print and log all headers for debugging
headers = dict(request.headers)
print(f"\n=== Agent Validator Middleware ===")
print(f"Path: {request.url.path}")
print(f"Method: {request.method}")
print("\nHeaders (case-sensitive):")
for name, value in request.headers.raw:
print(f" {name.decode()}: {value.decode()}")
print("\nParsed headers:")
print(f" X-Agent-Type: {headers.get('X-Agent-Type')}")
print(f" x-agent-type: {headers.get('x-agent-type')}")
print("===============================\n")
logger.info(f"Received headers: {headers}")
# Skip agent validation for the root endpoint and MCP endpoint
if request.url.path in ["/", "/mcp"]:
print("Skipping validation for root or MCP endpoint")
return await call_next(request)
# Skip validation if disabled in settings
if not settings.agent.AGENT_VALIDATION_ENABLED:
print("Agent validation is disabled in settings")
return await call_next(request)
# Get the X-Agent-Type header (case-insensitive)
agent_type = request.headers.get("X-Agent-Type") or request.headers.get("x-agent-type")
if not agent_type:
print("ERROR: X-Agent-Type header is missing")
raise HTTPException(
status_code=403,
detail="X-Agent-Type header is required"
)
# Check if the agent type is allowed
if agent_type.lower() not in [agent.lower() for agent in settings.agent.ALLOWED_AGENTS]:
print(f"ERROR: Invalid agent type: {agent_type}")
raise HTTPException(
status_code=403,
detail=f"Agent type '{agent_type}' is not allowed. Allowed agents: {', '.join(settings.agent.ALLOWED_AGENTS)}"
)
print(f"Agent validation passed for type: {agent_type}")
# If everything is valid, proceed with the request
return await call_next(request) |