diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 4b456710057..49434befd4e 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -7,12 +7,13 @@ import base64 import hashlib import json import os -from collections.abc import Awaitable, Callable, Generator +from collections.abc import Awaitable, Callable, Generator, Sequence from contextlib import AbstractAsyncContextManager from functools import partial from types import MappingProxyType from typing import Any, Final, TypeAlias, TypeVar +import anyio import httpx2 from httpx2._client import UseClientDefault from httpx2._types import AuthTypes @@ -38,6 +39,8 @@ from mcp.types import ( ListPromptsResult, ListResourcesResult, ListResourceTemplatesResult, + PaginatedRequestParams, + PaginatedResult, Prompt, ResourceTemplate, ServerNotification, @@ -49,7 +52,12 @@ from mcp.types import Tool as MCPTool from pydantic import AnyUrl from litellm._logging import verbose_logger -from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_TIMEOUT +from litellm.constants import ( + MCP_CLIENT_TIMEOUT, + MCP_NPM_CACHE_DIR, + MCP_TOOL_LISTING_MAX_PAGES, + MCP_TOOL_LISTING_TIMEOUT, +) from litellm.experimental_mcp_client.tools import list_tools_with_pagination from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response @@ -147,6 +155,8 @@ def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None: TSessionResult = TypeVar("TSessionResult") +_ListPage = TypeVar("_ListPage", bound=PaginatedResult) +_ListItem = TypeVar("_ListItem") class _MCPHTTPClient(httpx2.AsyncClient): @@ -793,6 +803,33 @@ class MCPClient: # Return a default error result instead of raising return self.error_tool_result(e) + async def _list_optional_pages( + self, + fetch_page: Callable[[PaginatedRequestParams | None], Awaitable[_ListPage]], + items_of: Callable[[_ListPage], Sequence[_ListItem]], + ) -> list[_ListItem]: # mutable-ok: existing list discovery API + items: Final[list[_ListItem]] = [] # mutable-ok: bounded iterative page accumulation + cursors: Final[set[str]] = set() # mutable-ok: constant-time detection of cursor cycles + cursor: str | None = None # rebind-ok: iterative traversal avoids recursion at the existing page cap + with anyio.fail_after(max(self.timeout, MCP_TOOL_LISTING_TIMEOUT)): + for page_index in range(MCP_TOOL_LISTING_MAX_PAGES): + try: + page = await fetch_page( # rebind-ok: each SDK page replaces the previous one + None if cursor is None else PaginatedRequestParams(cursor=cursor) + ) + except MCPError as error: + if page_index > 0 and error.error.code == METHOD_NOT_FOUND: + raise RuntimeError("MCP list operation became unavailable during pagination") from error + raise + items.extend(items_of(page)) + if not page.next_cursor: + return items + if page.next_cursor in cursors: + raise RuntimeError("MCP list pagination repeated a cursor") + cursors.add(page.next_cursor) + cursor = page.next_cursor + raise RuntimeError(f"MCP list pagination exceeded {MCP_TOOL_LISTING_MAX_PAGES} pages") + async def list_prompts(self, *, raise_on_error: bool = False) -> list[Prompt]: """List available prompts from the server.""" verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") @@ -802,7 +839,11 @@ class MCPClient: if capabilities is not None and capabilities.prompts is None: return ListPromptsResult(prompts=[]) try: - return await session.list_prompts() + return ListPromptsResult( + prompts=await self._list_optional_pages( + lambda params: session.list_prompts(params=params), lambda page: page.prompts + ) + ) except MCPError as error: if error.error.code != METHOD_NOT_FOUND: raise @@ -892,7 +933,11 @@ class MCPClient: if capabilities is not None and capabilities.resources is None: return ListResourcesResult(resources=[]) try: - return await session.list_resources() + return ListResourcesResult( + resources=await self._list_optional_pages( + lambda params: session.list_resources(params=params), lambda page: page.resources + ) + ) except MCPError as error: if error.error.code != METHOD_NOT_FOUND: raise @@ -941,7 +986,12 @@ class MCPClient: if capabilities is not None and capabilities.resources is None: return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload try: - return await session.list_resource_templates() + return ListResourceTemplatesResult( + resource_templates=await self._list_optional_pages( + lambda params: session.list_resource_templates(params=params), + lambda page: page.resource_templates, + ) + ) except MCPError as error: if error.error.code != METHOD_NOT_FOUND: raise diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 6c1b7946394..bf37b1be2e4 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -46,6 +46,8 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionDocumentObject, ChatCompletionNamedToolChoiceParam, + ChatCompletionRedactedThinkingBlock, + ChatCompletionThinkingBlock, ChatCompletionToolParam, OpenAIMessageContentListBlock, ) @@ -854,6 +856,8 @@ def _count_content_list( content_list: str | Iterable[ OpenAIMessageContentListBlock + | ChatCompletionThinkingBlock + | ChatCompletionRedactedThinkingBlock | AnthropicMessagesTextParam | AnthropicMessagesImageParam | AnthropicMessagesDocumentParam @@ -898,9 +902,9 @@ def _count_content_list( use_default_image_token_count, default_token_count, ) - elif c["type"] == "thinking": + elif c["type"] in ("thinking", "redacted_thinking"): # Claude extended thinking content block - # Count the thinking text and skip signature (opaque signature blob) + # Count the thinking text and skip the opaque blobs (signature, redacted data) thinking_text = str(c.get("thinking", "")) if thinking_text: num_tokens += count_function(thinking_text) @@ -920,7 +924,8 @@ def _count_content_list( raise ValueError( f"Invalid content item type: {content_type}. " f"Expected str or dict with 'type' field " - f"(text, image_url, image, document, file, tool_use, tool_result, thinking, tool_reference)." + f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, " + f"tool_reference)." ) return num_tokens except Exception as e: diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 4b698f1258d..6c20ef135ba 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -2036,6 +2036,15 @@ async def test_optional_discovery_preserves_cancellation(method: str) -> None: }, }, ) + if not (payload.params or {}).get("cursor"): + field: Final = { + "prompts/list": "prompts", + "resources/list": "resources", + "resources/templates/list": "resourceTemplates", + }[method] + return httpx2.Response( + 200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [], "nextCursor": "pending-page"}} + ) ready.set() await pending.wait() return httpx2.Response(202) @@ -2055,6 +2064,255 @@ async def test_optional_discovery_preserves_cancellation(method: str) -> None: await asyncio.wait_for(task, timeout=3) +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) +@pytest.mark.parametrize("session_id", (None, "pagination-session")) +@pytest.mark.parametrize("empty_middle", (False, True)) +async def test_optional_discovery_collects_all_pages(method: str, session_id: str | None, empty_middle: bool) -> None: + from mcp.types import Prompt, PromptArgument, Resource, ResourceTemplate + + field: Final = { + "prompts/list": "prompts", + "resources/list": "resources", + "resources/templates/list": "resourceTemplates", + }[method] + entries: Final = tuple( + { + "prompts/list": Prompt( + name=f"item-{index}", + description="prompt description", + arguments=[PromptArgument(name="query", required=True)], + ), + "resources/list": Resource( + name=f"item-{index}", + uri=f"test://item/{index}", + mime_type="text/plain", + description="resource description", + ), + "resources/templates/list": ResourceTemplate( + name=f"item-{index}", uri_template=f"test://item/{index}/{{query}}", mime_type="text/plain" + ), + }[method] + for index in range(5) + ) + + def respond(request: httpx2.Request) -> httpx2.Response: + if request.method == "GET": + return httpx2.Response(405) + if request.method == "DELETE": + return httpx2.Response(200) + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + if payload.method == "initialize": + return httpx2.Response( + 200, + headers={"mcp-session-id": session_id} if session_id else {}, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": payload.params["protocolVersion"], + "capabilities": {"prompts": {}, "resources": {}}, + "serverInfo": {"name": "paged", "version": "1"}, + }, + }, + ) + assert payload.method == method + assert request.headers.get("mcp-session-id") == session_id + cursor: Final = (payload.params or {}).get("cursor") + assert cursor in (None, "opaque:/second+page", "opaque:/last+page") + page: Final = ( + entries[:3] if cursor is None else (() if empty_middle and cursor == "opaque:/second+page" else entries[3:]) + ) + next_cursor: Final = ( + "opaque:/second+page" + if cursor is None + else "opaque:/last+page" + if empty_middle and cursor == "opaque:/second+page" + else "" + ) + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + field: [item.model_dump(mode="json", by_alias=True) for item in page], + "nextCursor": next_cursor, + }, + }, + ) + + responder: Final = Mock(side_effect=respond) + client: Final = _MockTransportClient(responder, server_url="https://example.com/mcp") + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + assert await operation(raise_on_error=True) == list(entries) + requests: Final = tuple( + _JSONRPC_MESSAGE_ADAPTER.validate_json(call.args[0].content) + for call in responder.call_args_list + if call.args[0].method == "POST" + ) + assert sum(isinstance(request, JSONRPCRequest) and request.method == "initialize" for request in requests) == 1 + assert tuple( + (request.params or {}).get("cursor") + for request in requests + if isinstance(request, JSONRPCRequest) and request.method == method + ) == ((None, "opaque:/second+page", "opaque:/last+page") if empty_middle else (None, "opaque:/second+page")) + assert sum(call.args[0].method == "DELETE" for call in responder.call_args_list) == (1 if session_id else 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) +@pytest.mark.parametrize( + "failure", ("repeat", "cycle", "cap", "method_not_found", "internal_error", "unauthorized", "deadline") +) +@pytest.mark.parametrize("strict", (False, True)) +async def test_optional_discovery_rejects_incomplete_walks( + method: str, failure: str, strict: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr(mcp_client_module, "MCP_TOOL_LISTING_MAX_PAGES", 3 if failure == "cycle" else 2, raising=False) + monkeypatch.setattr(mcp_client_module, "MCP_TOOL_LISTING_TIMEOUT", 0.05) + field: Final = { + "prompts/list": "prompts", + "resources/list": "resources", + "resources/templates/list": "resourceTemplates", + }[method] + entry: Final = { + "prompts/list": {"name": "first"}, + "resources/list": {"name": "first", "uri": "test://first"}, + "resources/templates/list": {"name": "first", "uriTemplate": "test://{name}"}, + }[method] + cancelled: Final = asyncio.Event() + + async def respond(request: httpx2.Request) -> httpx2.Response: + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + if payload.method == "initialize": + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": payload.params["protocolVersion"], + "capabilities": {"prompts": {}, "resources": {}}, + "serverInfo": {"name": "interrupted", "version": "1"}, + }, + }, + ) + assert payload.method == method + cursor: Final = (payload.params or {}).get("cursor") + if cursor is not None: + if failure == "deadline": + try: + await asyncio.Event().wait() + finally: + cancelled.set() + if failure == "unauthorized": + return httpx2.Response(401) + if failure in ("method_not_found", "internal_error"): + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "error": { + "code": -32601 if failure == "method_not_found" else -32603, + "message": "Later page unavailable", + }, + }, + ) + next_cursor: Final = ( + "private-cursor-2" if cursor == "private-cursor-1" and failure != "repeat" else "private-cursor-1" + ) + return httpx2.Response( + 200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [entry], "nextCursor": next_cursor}} + ) + + responder: Final = AsyncMock(side_effect=respond) + client: Final = _MockTransportClient(responder, server_url="https://example.com/mcp", timeout=0.2) + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + if strict: + error_type: Final = { + "internal_error": MCPError, + "unauthorized": httpx2.HTTPStatusError, + "deadline": TimeoutError, + }.get(failure, RuntimeError) + with pytest.raises(error_type): + await operation(raise_on_error=True) + else: + assert await operation() == [] + assert len( + tuple( + payload + for call in responder.call_args_list + if isinstance(payload := _JSONRPC_MESSAGE_ADAPTER.validate_json(call.args[0].content), JSONRPCRequest) + and payload.method == method + ) + ) == (3 if failure == "cycle" else 2) + assert "private-cursor" not in caplog.text + if failure == "deadline": + assert cancelled.is_set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) +async def test_optional_discovery_allows_exhaustion_at_page_cap(method: str, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(mcp_client_module, "MCP_TOOL_LISTING_MAX_PAGES", 2, raising=False) + field: Final = { + "prompts/list": "prompts", + "resources/list": "resources", + "resources/templates/list": "resourceTemplates", + }[method] + + def respond(request: httpx2.Request) -> httpx2.Response: + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + if payload.method == "initialize": + result: Final = { + "protocolVersion": payload.params["protocolVersion"], + "capabilities": {"prompts": {}, "resources": {}}, + "serverInfo": {"name": "empty-pages", "version": "1"}, + } + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + assert payload.method == method + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": {field: [], "nextCursor": None if (payload.params or {}).get("cursor") else "last-page"}, + }, + ) + + responder: Final = Mock(side_effect=respond) + client: Final = _MockTransportClient(responder, server_url="https://example.com/mcp") + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + assert await operation(raise_on_error=True) == [] + assert ( + sum( + isinstance(payload := _JSONRPC_MESSAGE_ADAPTER.validate_json(call.args[0].content), JSONRPCRequest) + and payload.method == method + for call in responder.call_args_list + ) + == 2 + ) + def test_client_import_before_proxy_credentials_succeeds_in_fresh_process(): import subprocess diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index ba3a6be609f..f19a8891609 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -1257,6 +1257,25 @@ def test_token_counter_with_thinking_content(): ), f"Expected minimal token count for empty thinking block, got {tokens_no_thinking}" + +def test_token_counter_with_redacted_thinking_content(): + """ + A replayed redacted_thinking block (Anthropic redacted reasoning, or the /v1/messages bridge's stand-in + for a reasoning item with no summary) counts zero tokens for its encrypted payload, like a thinking + block with no text. It used to raise, which made is_prompt_caching_valid_prompt return False and the + prompt_caching pre-call check stop pinning the deployment that held the cached prefix. + """ + model = "anthropic/claude-sonnet-4-5-20250929" + reply = {"type": "text", "text": "Draw from the box labeled Mixed, because that label must be wrong."} + redacted_block = {"type": "redacted_thinking", "data": "EqQBCkYIBRgCKkBjZ2xhc3M" * 30} + user_turn = {"role": "user", "content": [{"type": "text", "text": "Which box do you draw from?"}]} + follow_up = {"role": "user", "content": [{"type": "text", "text": "Restate that in one sentence."}]} + + without_block = [user_turn, {"role": "assistant", "content": [reply]}, follow_up] + with_block = [user_turn, {"role": "assistant", "content": [redacted_block, reply]}, follow_up] + + assert token_counter(model=model, messages=with_block) == token_counter(model=model, messages=without_block) + def test_token_counter_with_tool_reference_block(): """ Regression test: a message containing an Anthropic tool-search diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 9140ac61f1a..7418cf67e5f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13182,6 +13182,8 @@ class _DiscoveryUpstream: await self.release.wait() if self.outcome == "failure": return httpx2.Response(503) + if self.outcome == "paged_failure" and (payload.params or {}).get("cursor"): + return httpx2.Response(503) if self.outcome == "cancelled": raise asyncio.CancelledError() if self.outcome == "rejected": @@ -13196,7 +13198,12 @@ class _DiscoveryUpstream: }, "tools/list": {"tools": []}, }[payload.method] - return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + continuation: Final = ( + {"nextCursor": "last-page"} + if self.outcome in ("paged", "paged_failure") and not (payload.params or {}).get("cursor") + else {} + ) + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {**result, **continuation}}) @property def initializes(self) -> int: @@ -13262,6 +13269,29 @@ async def test_discovery_cache_empty_results_and_failures(kind: str, outcome: st assert upstream.initializes == 3 +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ("prompts", "resources", "templates")) +async def test_discovery_cache_retries_failed_pagination_before_caching_complete_list(kind: str) -> None: + manager: Final = MCPServerManager() + upstream: Final = _DiscoveryUpstream() + upstream.outcome = "paged_failure" + operation: Final = { + "prompts": manager.get_prompts_from_server, + "resources": manager.get_resources_from_server, + "templates": manager.get_resource_templates_from_server, + }[kind] + with _mcp_upstream(upstream.respond): + assert await operation(_discovery_server(), None) == [] + assert upstream.initializes == 1 + upstream.outcome = "paged" + recovered: Final = await operation(_discovery_server(), None) + assert [item.name for item in recovered] == ["discovery-example", "discovery-example"] + assert upstream.initializes == 2 + requests_after_recovery: Final = upstream.requests + assert await operation(_discovery_server(), None) == recovered + assert upstream.requests == requests_after_recovery + + @pytest.mark.asyncio async def test_discovery_cache_isolates_forwarded_credentials_and_shares_static_auth() -> None: import respx diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 849edc8c537..323abd68c59 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -210,6 +210,58 @@ async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is AUTO_CACHING_MODEL = "anthropic/claude-sonnet-4-5" +@pytest.mark.asyncio +async def test_replayed_redacted_thinking_block_still_records_and_pins(): + """ + A model that returns no reasoning summary (gpt-5.x through the /v1/messages bridge, Anthropic with + redacted reasoning) hands the client a `redacted_thinking` block, and the client replays it on every + later turn. The token count behind `is_prompt_caching_valid_prompt` raised on that block, the helper + swallowed it to False, and the check neither recorded the serving deployment nor pinned it, so the + conversation bounced across the group and paid a cache write on each deployment. + """ + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + model = "openai/gpt-5.6-sol" + deployments = _deployments(model, model, model) + messages = cast( + list[AllMessageValues], + [ + *_messages(word_count=3000), + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": "litellm_encrypted_reasoning:" + "Z" * 400}, + {"type": "text", "text": "Draw from the box labeled Mixed."}, + ], + }, + {"role": "user", "content": "Restate that in one sentence."}, + ], + ) + + assert is_prompt_caching_valid_prompt(model=model, messages=messages) is True + + await check.async_log_success_event( + kwargs={ + "standard_logging_object": { + "call_type": "anthropic_messages", + "model": model, + "messages": messages, + "model_id": "dep-2", + } + }, + response_obj=None, + start_time=None, + end_time=None, + ) + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + ) + + assert filtered == [deployments[1]] + + def _auto_caching_messages() -> list[AllMessageValues]: """A prompt over the model minimum that carries no client cache_control.""" return cast(