diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5b3d5bd60e2..cdb26acb658 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -5,7 +5,6 @@ LiteLLM MCP Server Routes import asyncio import contextlib - import traceback import uuid from datetime import datetime @@ -44,7 +43,10 @@ from litellm.proxy._experimental.mcp_server.utils import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.ip_address_utils import IPAddressUtils -from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + get_chain_id_from_headers, +) from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall @@ -331,6 +333,11 @@ if MCP_AVAILABLE: try: # Create a body date for logging body_data = {"name": name, "arguments": arguments} + # Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A) + chain_id = get_chain_id_from_headers(raw_headers) + if chain_id: + body_data["litellm_trace_id"] = chain_id + body_data["litellm_session_id"] = chain_id request = Request( scope={ @@ -884,6 +891,10 @@ if MCP_AVAILABLE: # This is intentionally minimal: only async_success_handler / post_call_failure_hook rules_obj = Rules() list_tools_call_id = str(uuid.uuid4()) + # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool) + effective_litellm_trace_id = litellm_trace_id or get_chain_id_from_headers( + raw_headers + ) spend_logs_metadata: Dict[str, Any] = { "mcp_operation": "list_tools", } @@ -896,7 +907,7 @@ if MCP_AVAILABLE: "model": "MCP: list_tools", "call_type": CallTypes.list_mcp_tools.value, "litellm_call_id": list_tools_call_id, - "litellm_trace_id": litellm_trace_id, + "litellm_trace_id": effective_litellm_trace_id, "metadata": { "spend_logs_metadata": spend_logs_metadata, }, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9861026ba3e..68f1ab114bf 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -10,12 +10,16 @@ import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.proxy._types import (AddTeamCallback, CommonProxyErrors, - LitellmDataForBackendLLMCall, - LitellmUserRoles, SpecialHeaders, - TeamCallbackMetadata, UserAPIKeyAuth) -from litellm.proxy.common_utils.http_parsing_utils import \ - _safe_get_request_headers +from litellm.proxy._types import ( + AddTeamCallback, + CommonProxyErrors, + LitellmDataForBackendLLMCall, + LitellmUserRoles, + SpecialHeaders, + TeamCallbackMetadata, + UserAPIKeyAuth, +) +from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers # Cache special headers as a frozenset for O(1) lookup performance _SPECIAL_HEADERS_CACHE = frozenset( @@ -24,9 +28,12 @@ _SPECIAL_HEADERS_CACHE = frozenset( from litellm.router import Router from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS from litellm.types.services import ServiceTypes -from litellm.types.utils import (LlmProviders, ProviderSpecificHeader, - StandardLoggingUserAPIKeyMetadata, - SupportedCacheControls) +from litellm.types.utils import ( + LlmProviders, + ProviderSpecificHeader, + StandardLoggingUserAPIKeyMetadata, + SupportedCacheControls, +) service_logger_obj = ServiceLogging() # used for tracking latency on OTEL @@ -82,6 +89,25 @@ def _get_metadata_variable_name(request: Request) -> str: return "metadata" +def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str]: + """ + Extract chain id for call chaining from request headers. + + x-litellm-trace-id and x-litellm-session-id are interchangeable; when both + are present, x-litellm-trace-id takes precedence. Header keys are matched + case-insensitively so this works with raw header dicts from any transport. + + Used by MCP (and other paths that have raw_headers but no Request) to set + litellm_trace_id/litellm_session_id for spend logs and logging consistency. + """ + if not headers: + return None + normalized = {k.lower(): v for k, v in headers.items() if isinstance(k, str)} + return normalized.get("x-litellm-trace-id") or normalized.get( + "x-litellm-session-id" + ) + + def safe_add_api_version_from_query_params(data: dict, request: Request): try: if hasattr(request, "query_params"): @@ -663,7 +689,8 @@ class LiteLLMProxyRequestSetup: return data from litellm.proxy._types import ( LiteLLM_ManagementEndpoint_MetadataFields, - LiteLLM_ManagementEndpoint_MetadataFields_Premium) + LiteLLM_ManagementEndpoint_MetadataFields_Premium, + ) # ignore any special fields added_metadata = {} @@ -832,8 +859,7 @@ async def add_litellm_data_to_request( # noqa: PLR0915 """ from litellm.proxy.proxy_server import llm_router, premium_user - from litellm.types.proxy.litellm_pre_call_utils import (RedactedDict, - SecretFields) + from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields _raw_headers: Dict[str, str] = RedactedDict(_safe_get_request_headers(request)) @@ -1508,8 +1534,7 @@ async def move_guardrails_to_metadata( # Only check policy engine if no local config (avoid import + registry lookup) if not (has_key_config or has_team_config or has_request_config): - from litellm.proxy.policy_engine.policy_registry import \ - get_policy_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry if not get_policy_registry().is_initialized(): # Nothing configured anywhere - clean up request body fields and return @@ -1573,16 +1598,14 @@ async def move_guardrails_to_metadata( def _is_policy_version_id(s: str) -> bool: """Return True if string is a policy version ID (starts with policy_ prefix).""" - from litellm.proxy.policy_engine.policy_registry import \ - POLICY_VERSION_ID_PREFIX + from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX return isinstance(s, str) and s.startswith(POLICY_VERSION_ID_PREFIX) def _extract_policy_id(s: str) -> Optional[str]: """Extract raw UUID from policy_ string, or None if not a valid version ID.""" - from litellm.proxy.policy_engine.policy_registry import \ - POLICY_VERSION_ID_PREFIX + from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX if not _is_policy_version_id(s): return None @@ -1603,9 +1626,10 @@ def _match_and_track_policies( """ from litellm._logging import verbose_proxy_logger from litellm.proxy.common_utils.callback_utils import ( - add_policy_sources_to_metadata, add_policy_to_applied_policies_header) - from litellm.proxy.policy_engine.attachment_registry import \ - get_attachment_registry + add_policy_sources_to_metadata, + add_policy_to_applied_policies_header, + ) + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher # Get matching policies via attachments (with match reasons for attribution) @@ -1750,8 +1774,7 @@ async def add_guardrails_from_policy_engine( user_api_key_dict: The user's API key authentication info """ from litellm._logging import verbose_proxy_logger - from litellm.proxy.common_utils.http_parsing_utils import \ - get_tags_from_request_body + from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.types.proxy.policy_engine import PolicyMatchContext