From 33531649c305ad48af390fcc93581fa2d3bb16f6 Mon Sep 17 00:00:00 2001
From: yassin <yassin@berri.ai>
Date: Thu, 17 Sep 2026 23:20:47 +0000
Subject: [PATCH] perf(mcp): count gateway session groups with Counter and pin
 the oversized initialize peek invariant

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../proxy/_experimental/mcp_server/server.py  |  5 +++--
 .../mcp_server/test_mcp_server.py             | 22 +++++++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index 9c8ad2f4613..9adc203b092 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -14,6 +14,7 @@ import time
 import traceback
 import types
 import uuid
+from collections import Counter
 from collections.abc import AsyncIterator, Callable, Mapping, Sequence
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol
@@ -3837,10 +3838,10 @@ if MCP_AVAILABLE:
         sessions: Sequence[MCPGatewaySession],
         label_for: Callable[[MCPGatewaySession], str | None],
     ) -> tuple[MCPGatewaySessionGroupCount, ...]:
-        labels: Final = tuple(label_for(session) for session in sessions)
+        counts: Final = Counter(label_for(session) for session in sessions)
         return tuple(
             sorted(
-                (MCPGatewaySessionGroupCount(label=label, count=labels.count(label)) for label in frozenset(labels)),
+                (MCPGatewaySessionGroupCount(label=label, count=count) for label, count in counts.items()),
                 key=lambda group: (-group.count, group.label is None, group.label or ""),
             )
         )
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
index bbc36991e21..ef424255f04 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
@@ -2658,6 +2658,28 @@ def test_extract_initialize_client_info_returns_none_without_client_info(body):
     assert mcp_server._extract_initialize_client_info(body) is None
 
 
+def test_oversized_initialize_peek_neither_routes_stateful_nor_attributes_client():
+    """The routing sniff and the clientInfo parse read the same capped peek, so
+    an initialize larger than the peek can never become a tracked session that
+    then reports an unknown client."""
+    try:
+        from litellm.proxy._experimental.mcp_server import server as mcp_server
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    padding = "x" * (mcp_server._MCP_ROUTING_PEEK_MAX_BYTES + 512)
+    full_body = (
+        b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18",'
+        b'"capabilities":{"experimental":{"pad":{"value":"' + padding.encode() + b'"}}},'
+        b'"clientInfo":{"name":"claude-code","version":"1.0.0"}}}'
+    )
+    peeked = full_body[: mcp_server._MCP_ROUTING_PEEK_MAX_BYTES]
+
+    assert mcp_server._extract_initialize_client_info(full_body) is not None
+    assert mcp_server._is_initialize_request(peeked) is False
+    assert mcp_server._extract_initialize_client_info(peeked) is None
+
+
 @pytest.mark.asyncio
 async def test_initialize_request_records_client_name_in_gateway_sessions_report():
     """The real initialize body's clientInfo is attributed to the session the