Add correlation ID system for unified request tracking

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

Author: shoummu1

.env.example

@@ -653,6 +653,17 @@ LOG_MAX_SIZE_MB=1
 LOG_BACKUP_COUNT=5
 LOG_BUFFER_SIZE_MB=1.0
 
+# Correlation ID / Request Tracking
+# Enable automatic correlation ID tracking for unified request tracing
+# Options: true (default), false
+CORRELATION_ID_ENABLED=true
+# HTTP header name for correlation ID (default: X-Correlation-ID)
+CORRELATION_ID_HEADER=X-Correlation-ID
+# Preserve incoming correlation IDs from clients (default: true)
+CORRELATION_ID_PRESERVE=true
+# Include correlation ID in HTTP response headers (default: true)
+CORRELATION_ID_RESPONSE_HEADER=true
+
 # Transport Protocol Configuration
 # Options: all (default), sse, streamablehttp, http
 # - all: Enable all transport protocols

mcpgateway/auth.py

@@ -27,10 +27,62 @@
 from mcpgateway.db import EmailUser, SessionLocal
 from mcpgateway.plugins.framework import get_plugin_manager, GlobalContext, HttpAuthResolveUserPayload, HttpHeaderPayload, HttpHookType, PluginViolationError
 from mcpgateway.services.team_management_service import TeamManagementService  # pylint: disable=import-outside-toplevel
+from mcpgateway.utils.correlation_id import get_correlation_id
 from mcpgateway.utils.verify_credentials import verify_jwt_token
 
 # Security scheme
-bearer_scheme = HTTPBearer(auto_error=False)
+security = HTTPBearer(auto_error=False)
+
+
+def _log_auth_event(
+    logger: logging.Logger,
+    message: str,
+    level: int = logging.INFO,
+    user_id: Optional[str] = None,
+    auth_method: Optional[str] = None,
+    auth_success: bool = False,
+    security_event: Optional[str] = None,
+    security_severity: str = "low",
+    **extra_context
+) -> None:
+    """Log authentication event with structured context and request_id.
+
+    This helper creates structured log records that include request_id from the
+    correlation ID context, enabling end-to-end tracing of authentication flows.
+
+    Args:
+        logger: Logger instance to use
+        message: Log message
+        level: Log level (default: INFO)
+        user_id: User identifier
+        auth_method: Authentication method used (jwt, api_token, etc.)
+        auth_success: Whether authentication succeeded
+        security_event: Type of security event (authentication, authorization, etc.)
+        security_severity: Severity level (low, medium, high, critical)
+        **extra_context: Additional context fields
+    """
+    # Get request_id from correlation ID context
+    request_id = get_correlation_id()
+
+    # Build structured log record
+    extra = {
+        'request_id': request_id,
+        'entity_type': 'auth',
+        'auth_success': auth_success,
+        'security_event': security_event or 'authentication',
+        'security_severity': security_severity,
+    }
+
+    if user_id:
+        extra['user_id'] = user_id
+    if auth_method:
+        extra['auth_method'] = auth_method
+
+    # Add any additional context
+    extra.update(extra_context)
+
+    # Log with structured context
+    logger.log(level, message, extra=extra)
 
 
 def get_db() -> Generator[Session, Never, None]:
@@ -169,12 +221,15 @@ async def get_current_user(
             if request and hasattr(request, "headers"):
                 headers = dict(request.headers)
 
-            # Get request ID from request state (set by middleware) or generate new one
-            request_id = None
-            if request and hasattr(request, "state") and hasattr(request.state, "request_id"):
-                request_id = request.state.request_id
-            else:
-                request_id = uuid.uuid4().hex
+            # Get request ID from correlation ID context (set by CorrelationIDMiddleware)
+            request_id = get_correlation_id()
+            if not request_id:
+                # Fallback chain for safety
+                if request and hasattr(request, "state") and hasattr(request.state, "request_id"):
+                    request_id = request.state.request_id
+                else:
+                    request_id = uuid.uuid4().hex
+                    logger.debug(f"Generated fallback request ID in get_current_user: {request_id}")
 
             # Create global context
             global_context = GlobalContext(

mcpgateway/config.py

@@ -747,6 +747,12 @@ def _parse_allowed_origins(cls, v: Any) -> Set[str]:
     # Enable span events
     observability_events_enabled: bool = Field(default=True, description="Enable event logging within spans")
 
+    # Correlation ID Settings
+    correlation_id_enabled: bool = Field(default=True, description="Enable automatic correlation ID tracking for requests")
+    correlation_id_header: str = Field(default="X-Correlation-ID", description="HTTP header name for correlation ID")
+    correlation_id_preserve: bool = Field(default=True, description="Preserve correlation IDs from incoming requests")
+    correlation_id_response_header: bool = Field(default=True, description="Include correlation ID in response headers")
+
     @field_validator("log_level", mode="before")
     @classmethod
     def validate_log_level(cls, v: str) -> str:

mcpgateway/main.py

@@ -70,6 +70,7 @@
 from mcpgateway.db import refresh_slugs_on_startup, SessionLocal
 from mcpgateway.db import Tool as DbTool
 from mcpgateway.handlers.sampling import SamplingHandler
+from mcpgateway.middleware.correlation_id import CorrelationIDMiddleware
 from mcpgateway.middleware.http_auth_middleware import HttpAuthMiddleware
 from mcpgateway.middleware.protocol_version import MCPProtocolVersionMiddleware
 from mcpgateway.middleware.rbac import get_current_user_with_permissions, require_permission
@@ -1065,6 +1066,13 @@ async def _call_streamable_http(self, scope, receive, send):
 # Add HTTP authentication hook middleware for plugins (before auth dependencies)
 if plugin_manager:
     app.add_middleware(HttpAuthMiddleware, plugin_manager=plugin_manager)
+    logger.info("🔌 HTTP authentication hooks enabled for plugins")
+
+# Add correlation ID middleware if enabled
+# Note: Registered AFTER HttpAuthMiddleware so it executes FIRST (middleware runs in LIFO order)
+if settings.correlation_id_enabled:
+    app.add_middleware(CorrelationIDMiddleware)
+    logger.info(f"✅ Correlation ID tracking enabled (header: {settings.correlation_id_header})")
 
 # Add custom DocsAuthMiddleware
 app.add_middleware(DocsAuthMiddleware)

mcpgateway/middleware/correlation_id.py

@@ -0,0 +1,121 @@
+# -*- coding: utf-8 -*-
+"""Location: ./mcpgateway/middleware/correlation_id.py
+Copyright 2025
+SPDX-License-Identifier: Apache-2.0
+Authors: MCP Gateway Contributors
+
+Correlation ID (Request ID) Middleware.
+
+This middleware handles X-Correlation-ID HTTP headers and maps them to the internal
+request_id used throughout the system for unified request tracing.
+
+Key concept: HTTP X-Correlation-ID header → Internal request_id field (single ID for entire request flow)
+
+The middleware automatically extracts or generates request IDs for every HTTP request,
+stores them in context variables for async-safe propagation across services, and
+injects them back into response headers for client-side correlation.
+
+This enables end-to-end tracing: HTTP → Middleware → Services → Plugins → Logs (all with same request_id)
+"""
+
+# Standard
+import logging
+from typing import Callable
+
+# Third-Party
+from fastapi import Request, Response
+from starlette.middleware.base import BaseHTTPMiddleware
+
+# First-Party
+from mcpgateway.config import settings
+from mcpgateway.utils.correlation_id import (
+    clear_correlation_id,
+    extract_correlation_id_from_headers,
+    generate_correlation_id,
+    set_correlation_id,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class CorrelationIDMiddleware(BaseHTTPMiddleware):
+    """Middleware for automatic request ID (correlation ID) handling.
+
+    This middleware:
+    1. Extracts request ID from X-Correlation-ID header in incoming requests
+    2. Generates a new UUID if no correlation ID is present
+    3. Stores the ID in context variables for the request lifecycle (used as request_id throughout system)
+    4. Injects the request ID into X-Correlation-ID response header
+    5. Cleans up context after request completion
+
+    The request ID extracted/generated here becomes the unified request_id used in:
+    - All log entries (request_id field)
+    - GlobalContext.request_id (when plugins execute)
+    - Service method calls for tracing
+    - Database queries for request tracking
+
+    Configuration is controlled via settings:
+    - correlation_id_enabled: Enable/disable the middleware
+    - correlation_id_header: Header name to use (default: X-Correlation-ID)
+    - correlation_id_preserve: Whether to preserve incoming IDs (default: True)
+    - correlation_id_response_header: Whether to add ID to responses (default: True)
+    """
+
+    def __init__(self, app):
+        """Initialize the correlation ID (request ID) middleware.
+
+        Args:
+            app: The FastAPI application instance
+        """
+        super().__init__(app)
+        self.header_name = getattr(settings, 'correlation_id_header', 'X-Correlation-ID')
+        self.preserve_incoming = getattr(settings, 'correlation_id_preserve', True)
+        self.add_to_response = getattr(settings, 'correlation_id_response_header', True)
+
+    async def dispatch(self, request: Request, call_next: Callable) -> Response:
+        """Process the request and manage request ID (correlation ID) lifecycle.
+
+        Extracts or generates a request ID, stores it in context variables for use throughout
+        the request lifecycle (becomes request_id in logs, services, plugins), and injects
+        it back into the X-Correlation-ID response header.
+
+        Args:
+            request: The incoming HTTP request
+            call_next: The next middleware or route handler
+
+        Returns:
+            Response: The HTTP response with correlation ID header added
+        """
+        # Extract correlation ID from incoming request headers
+        correlation_id = None
+        if self.preserve_incoming:
+            correlation_id = extract_correlation_id_from_headers(
+                dict(request.headers),
+                self.header_name
+            )
+
+        # Generate new correlation ID if none was provided
+        if not correlation_id:
+            correlation_id = generate_correlation_id()
+            logger.debug(f"Generated new correlation ID: {correlation_id}")
+        else:
+            logger.debug(f"Using client-provided correlation ID: {correlation_id}")
+
+        # Store correlation ID in context variable for this request
+        # This makes it available to all downstream code (auth, services, plugins, logs)
+        set_correlation_id(correlation_id)
+
+        try:
+            # Process the request
+            response = await call_next(request)
+
+            # Add correlation ID to response headers if enabled
+            if self.add_to_response:
+                response.headers[self.header_name] = correlation_id
+
+            return response
+
+        finally:
+            # Clean up context after request completes
+            # Note: ContextVar automatically cleans up, but explicit cleanup is good practice
+            clear_correlation_id()

mcpgateway/middleware/http_auth_middleware.py

@@ -17,6 +17,7 @@
 
 # First-Party
 from mcpgateway.plugins.framework import GlobalContext, HttpHeaderPayload, HttpHookType, HttpPostRequestPayload, HttpPreRequestPayload, PluginManager
+from mcpgateway.utils.correlation_id import generate_correlation_id, get_correlation_id
 
 logger = logging.getLogger(__name__)
 
@@ -60,9 +61,14 @@ async def dispatch(self, request: Request, call_next):
         if not self.plugin_manager:
             return await call_next(request)
 
-        # Generate request ID for tracing and store in request state
-        # This ensures all hooks and downstream code see the same request ID
-        request_id = uuid.uuid4().hex
+        # Use correlation ID from CorrelationIDMiddleware if available
+        # This ensures all hooks and downstream code see the same unified request ID
+        request_id = get_correlation_id()
+        if not request_id:
+            # Fallback if correlation ID middleware is disabled
+            request_id = generate_correlation_id()
+            logger.debug(f"Correlation ID not found, generated fallback: {request_id}")
+
         request.state.request_id = request_id
 
         # Create global context for hooks

mcpgateway/middleware/request_logging_middleware.py

@@ -23,6 +23,7 @@
 
 # First-Party
 from mcpgateway.services.logging_service import LoggingService
+from mcpgateway.utils.correlation_id import get_correlation_id
 
 # Initialize logging service first
 logging_service = LoggingService()
@@ -171,12 +172,16 @@ async def dispatch(self, request: Request, call_next: Callable):
             # Mask sensitive headers
             masked_headers = mask_sensitive_headers(dict(request.headers))
 
+            # Get correlation ID for request tracking
+            request_id = get_correlation_id()
+
             logger.log(
                 log_level,
                 f"📩 Incoming request: {request.method} {request.url.path}\n"
                 f"Query params: {dict(request.query_params)}\n"
                 f"Headers: {masked_headers}\n"
                 f"Body: {payload_str}{'... [truncated]' if truncated else ''}",
+                extra={"request_id": request_id},
             )
 
         except Exception as e:

mcpgateway/observability.py

@@ -393,6 +393,24 @@ def create_span(name: str, attributes: Optional[Dict[str, Any]] = None) -> Any:
         # Return a no-op context manager if tracing is not configured or available
         return nullcontext()
 
+    # Auto-inject correlation ID into all spans for request tracing
+    try:
+        # Import here to avoid circular dependency
+        from mcpgateway.utils.correlation_id import get_correlation_id
+
+        correlation_id = get_correlation_id()
+        if correlation_id:
+            if attributes is None:
+                attributes = {}
+            # Add correlation ID if not already present
+            if "correlation_id" not in attributes:
+                attributes["correlation_id"] = correlation_id
+            if "request_id" not in attributes:
+                attributes["request_id"] = correlation_id  # Alias for compatibility
+    except ImportError:
+        # Correlation ID module not available, continue without it
+        pass
+
     # Start span and return the context manager
     span_context = _TRACER.start_as_current_span(name)
 

mcpgateway/services/a2a_service.py

@@ -28,6 +28,7 @@
 from mcpgateway.services.logging_service import LoggingService
 from mcpgateway.services.team_management_service import TeamManagementService
 from mcpgateway.services.tool_service import ToolService
+from mcpgateway.utils.correlation_id import get_correlation_id
 from mcpgateway.utils.create_slug import slugify
 from mcpgateway.utils.services_auth import encode_auth  # ,decode_auth