# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/middleware/token_usage_middleware.py
Copyright 2025
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti
Token Usage Logging Middleware.
This middleware logs API token usage for analytics and security monitoring.
It records each request made with an API token, including endpoint, method,
response time, and status code.
Note: Implemented as raw ASGI middleware (not BaseHTTPMiddleware) to avoid
response body buffering issues with streaming responses.
Examples:
>>> from mcpgateway.middleware.token_usage_middleware import TokenUsageMiddleware # doctest: +SKIP
>>> app.add_middleware(TokenUsageMiddleware) # doctest: +SKIP
"""
# Standard
import logging
import time
# Third-Party
from starlette.datastructures import Headers
from starlette.requests import Request
from starlette.types import ASGIApp, Receive, Scope, Send
# First-Party
from mcpgateway.db import fresh_db_session
from mcpgateway.middleware.path_filter import should_skip_auth_context
from mcpgateway.services.token_catalog_service import TokenCatalogService
from mcpgateway.utils.verify_credentials import verify_jwt_token_cached
logger = logging.getLogger(__name__)
class TokenUsageMiddleware:
"""Raw ASGI middleware for logging API token usage.
This middleware tracks when API tokens are used, recording details like:
- Endpoint accessed
- HTTP method
- Response status code
- Response time
- Client IP and user agent
This data is used for security auditing, usage analytics, and detecting
anomalous token usage patterns.
Note:
Only logs usage for requests authenticated with API tokens (identified
by request.state.auth_method == "api_token").
Implemented as raw ASGI middleware to avoid BaseHTTPMiddleware issues:
- BaseHTTPMiddleware buffers entire response bodies (problematic for streaming)
- Raw ASGI middleware streams responses efficiently
"""
def __init__(self, app: ASGIApp) -> None:
"""Initialize middleware.
Args:
app: ASGI application to wrap
"""
self.app = app
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
"""Process ASGI request.
Args:
scope: ASGI scope dict
receive: Receive callable
send: Send callable
"""
# Only process HTTP requests
if scope["type"] != "http":
await self.app(scope, receive, send)
return
# Skip health checks and static files
path = scope.get("path", "")
if should_skip_auth_context(path):
await self.app(scope, receive, send)
return
# Record start time
start_time = time.time()
# Capture response status
status_code = 200 # Default
async def send_wrapper(message: dict) -> None:
"""Wrap send to capture response status.
Args:
message: ASGI message dict containing response data
"""
nonlocal status_code
if message["type"] == "http.response.start":
status_code = message["status"]
await send(message)
# Process request
await self.app(scope, receive, send_wrapper)
# Calculate response time
response_time_ms = round((time.time() - start_time) * 1000)
# Only log if this was an API token request
state = scope.get("state", {})
auth_method = state.get("auth_method") if state else None
if auth_method != "api_token":
return
# Extract token information from scope state
jti = state.get("jti") if state else None
user = state.get("user") if state else None
user_email = getattr(user, "email", None) if user else None
if not user_email:
user_email = state.get("user_email") if state else None
# If we don't have JTI or email, try to decode the token
if not jti or not user_email:
try:
# Get token from Authorization header
headers = Headers(scope=scope)
auth_header = headers.get("authorization")
if not auth_header or not auth_header.startswith("Bearer "):
return
token = auth_header.replace("Bearer ", "")
# Decode token to get JTI and user email
# Note: We need to create a minimal Request-like object
request = Request(scope, receive)
try:
payload = await verify_jwt_token_cached(token, request)
jti = jti or payload.get("jti")
user_email = user_email or payload.get("sub") or payload.get("email")
except Exception as decode_error:
logger.debug(f"Failed to decode token for usage logging: {decode_error}")
return
except Exception as e:
logger.debug(f"Error extracting token information: {e}")
return
if not jti or not user_email:
logger.debug("Missing JTI or user_email for token usage logging")
return
# Log token usage
try:
with fresh_db_session() as db:
token_service = TokenCatalogService(db)
# Get client IP
client = scope.get("client")
ip_address = client[0] if client else None
# Get user agent
headers = Headers(scope=scope)
user_agent = headers.get("user-agent")
# Log usage
await token_service.log_token_usage(
jti=jti,
user_email=user_email,
endpoint=path,
method=scope.get("method", "GET"),
ip_address=ip_address,
user_agent=user_agent,
status_code=status_code,
response_time_ms=response_time_ms,
blocked=False,
block_reason=None,
)
except Exception as e:
logger.debug(f"Failed to log token usage: {e}")
