apply ruff formatting

This commit is contained in:
Taylor Wilsdon
2025-12-13 13:49:28 -08:00
parent 1d80a24ca4
commit 6b8352a354
50 changed files with 4010 additions and 2842 deletions

View File

@@ -26,24 +26,26 @@ class MCPSessionMiddleware(BaseHTTPMiddleware):
Middleware that extracts session information from requests and makes it
available to MCP tool functions via context variables.
"""
async def dispatch(self, request: Request, call_next: Callable) -> Any:
"""Process request and set session context."""
logger.debug(f"MCPSessionMiddleware processing request: {request.method} {request.url.path}")
logger.debug(
f"MCPSessionMiddleware processing request: {request.method} {request.url.path}"
)
# Skip non-MCP paths
if not request.url.path.startswith("/mcp"):
logger.debug(f"Skipping non-MCP path: {request.url.path}")
return await call_next(request)
session_context = None
try:
# Extract session information
headers = dict(request.headers)
session_id = extract_session_from_headers(headers)
# Try to get OAuth 2.1 auth context from FastMCP
auth_context = None
user_email = None
@@ -52,28 +54,33 @@ class MCPSessionMiddleware(BaseHTTPMiddleware):
if hasattr(request.state, "auth"):
auth_context = request.state.auth
# Extract user email from auth claims if available
if hasattr(auth_context, 'claims') and auth_context.claims:
user_email = auth_context.claims.get('email')
if hasattr(auth_context, "claims") and auth_context.claims:
user_email = auth_context.claims.get("email")
# Check for FastMCP session ID (from streamable HTTP transport)
if hasattr(request.state, "session_id"):
mcp_session_id = request.state.session_id
logger.debug(f"Found FastMCP session ID: {mcp_session_id}")
# Also check Authorization header for bearer tokens
auth_header = headers.get("authorization")
if auth_header and auth_header.lower().startswith("bearer ") and not user_email:
if (
auth_header
and auth_header.lower().startswith("bearer ")
and not user_email
):
try:
import jwt
token = auth_header[7:] # Remove "Bearer " prefix
# Decode without verification to extract email
claims = jwt.decode(token, options={"verify_signature": False})
user_email = claims.get('email')
user_email = claims.get("email")
if user_email:
logger.debug(f"Extracted user email from JWT: {user_email}")
except Exception:
pass
# Build session context
if session_id or auth_context or user_email or mcp_session_id:
# Create session ID hierarchy: explicit session_id > Google user session > FastMCP session
@@ -82,10 +89,11 @@ class MCPSessionMiddleware(BaseHTTPMiddleware):
effective_session_id = f"google_{user_email}"
elif not effective_session_id and mcp_session_id:
effective_session_id = mcp_session_id
session_context = SessionContext(
session_id=effective_session_id,
user_id=user_email or (auth_context.user_id if auth_context else None),
user_id=user_email
or (auth_context.user_id if auth_context else None),
auth_context=auth_context,
request=request,
metadata={
@@ -93,19 +101,19 @@ class MCPSessionMiddleware(BaseHTTPMiddleware):
"method": request.method,
"user_email": user_email,
"mcp_session_id": mcp_session_id,
}
},
)
logger.debug(
f"MCP request with session: session_id={session_context.session_id}, "
f"user_id={session_context.user_id}, path={request.url.path}"
)
# Process request with session context
with SessionContextManager(session_context):
response = await call_next(request)
return response
except Exception as e:
logger.error(f"Error in MCP session middleware: {e}")
# Continue without session context