scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

Merge remote-tracking branch 'origin/main' into litellm_add_edenai_provider

Browse source
This commit is contained in:
yassin 2026-09-21 19:40:48 +00:00
commit 889b8fb220
6 changed files with 423 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

60
litellm/experimental_mcp_client/client.py
View file

@ -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

11
litellm/litellm_core_utils/token_counter.py
View file

@ -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:

258
tests/test_litellm/experimental_mcp_client/test_mcp_client.py
View file

@ -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

19
tests/test_litellm/litellm_core_utils/test_token_counter.py
View file

@ -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

32
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -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

52
tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
View file

@ -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(