scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 131

fix(jwt/mcp): warn on unscoped JWT fallback; route agent permission lookup through shared helper

Browse source 
- _build_decode_kwargs no longer suppresses the unscoped-fallback warning
  when LiteLLM_JWTAuth.issuers is set: tokens whose iss does not match
  any configured issuer still fall through to the global path, and that
  fallback is itself unscoped when JWT_AUDIENCE/JWT_ISSUER are absent.

- _get_agent_object_permission now caches the agent_id ->
  object_permission_id mapping and delegates the permission lookup to
  the shared get_object_permission helper, so the agent path reuses the
  same cache entries as the org / team / key paths.
This commit is contained in:
mateo-berri 2026-05-21 19:14:18 +00:00
parent ad8c9c2aa2
commit 1b0b75ccd7
No known key found for this signature in database
4 changed files with 135 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

74
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -1289,12 +1289,17 @@ class MCPRequestHandler:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
):
"""
Get agent object_permission, using user_api_key_cache to avoid DB hits on every request.
Caches both positive results and the absence of an object_permission so that agents
with no MCP permissions configured do not trigger a DB query on every request.
Get agent object_permission via the established ``get_object_permission``
helper. Caches the ``agent_id -> object_permission_id`` mapping so we
avoid re-reading the agent row on every request, and reuses the shared
``object_permission_id`` cache populated by the org / team / key paths.
"""
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
from litellm.proxy.auth.auth_checks import get_object_permission
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if not user_api_key_auth or not user_api_key_auth.agent_id:
return None
@ -1304,42 +1309,41 @@ class MCPRequestHandler:
return None
agent_id = user_api_key_auth.agent_id
cache_key = f"agent_object_permission:{agent_id}"
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
cache_key = f"agent_object_permission_id:{agent_id}"
try:
cached = await user_api_key_cache.async_get_cache(key=cache_key)
if cached is not None:
if cached == MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL:
return None
# Redis deserialises to a plain dict; reconstruct the Pydantic model
# so callers can access .mcp_servers / .mcp_tool_permissions as attrs.
if isinstance(cached, dict):
return LiteLLM_ObjectPermissionTable(**cached)
return cached
agent_row = await prisma_client.db.litellm_agentstable.find_unique(
where={"agent_id": agent_id},
include={"object_permission": True},
object_permission_id: Optional[str] = (
await user_api_key_cache.async_get_cache(key=cache_key)
)
if agent_row is None or agent_row.object_permission is None:
await user_api_key_cache.async_set_cache(
key=cache_key,
value=MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
if object_permission_id == MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL:
return None
obj_perm = LiteLLM_ObjectPermissionTable(
**agent_row.object_permission.dict()
if object_permission_id is None:
agent_row = await prisma_client.db.litellm_agentstable.find_unique(
where={"agent_id": agent_id},
)
object_permission_id = (
getattr(agent_row, "object_permission_id", None)
if agent_row is not None
else None
)
await user_api_key_cache.async_set_cache(
key=cache_key,
value=object_permission_id
or MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
if not object_permission_id:
return None
return await get_object_permission(
object_permission_id=object_permission_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
await user_api_key_cache.async_set_cache(
key=cache_key,
value=obj_perm.dict(),
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
return obj_perm
except Exception as e:
verbose_logger.warning(f"Failed to get agent object permission: {str(e)}")
return None

17
litellm/proxy/auth/handle_jwt.py
View file

@ -843,7 +843,7 @@ class JWTHandler:
_unscoped_jwt_warning_emitted = False
@classmethod
def _build_decode_kwargs(cls, has_issuer_config: bool = False) -> dict:
def _build_decode_kwargs(cls) -> dict:
"""Build the audience/issuer/options kwargs for ``jwt.decode``.
Setting ``JWT_AUDIENCE`` (and optionally ``JWT_ISSUER``) turns on the
@ -852,10 +852,10 @@ class JWTHandler:
When both are unset PyJWT only checks the signature and expiry, which
is preserved for backward compatibility but logged once as a warning.
``has_issuer_config`` suppresses the warning when the caller has
configured per-issuer scoping via ``LiteLLM_JWTAuth.issuers``: this
global path is then just the fallback for tokens whose ``iss`` did not
match any configured issuer, not the only scoping mechanism.
The warning fires even in mixed deployments that also configure
``LiteLLM_JWTAuth.issuers``: tokens whose ``iss`` does not match any
configured issuer fall through to this global path, and if env-var
scoping is absent that fallback is itself unscoped.
"""
audience = os.getenv("JWT_AUDIENCE")
issuer = os.getenv("JWT_ISSUER")
@ -863,7 +863,6 @@ class JWTHandler:
if (
audience is None
and issuer is None
and not has_issuer_config
and not cls._unscoped_jwt_warning_emitted
):
verbose_proxy_logger.warning(
@ -1057,11 +1056,7 @@ class JWTHandler:
kid=kid,
)
decode_kwargs = self._build_decode_kwargs(
has_issuer_config=bool(
getattr(getattr(self, "litellm_jwtauth", None), "issuers", None)
)
)
decode_kwargs = self._build_decode_kwargs()
public_key = await self.get_public_key(kid=kid)

83
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

@ -3202,6 +3202,89 @@ class TestAgentMCPPermissions:
)
assert sorted(result) == ["tool_a", "tool_b"]
async def test_get_agent_object_permission_uses_shared_helper(self):
"""``_get_agent_object_permission`` must resolve the agent's
``object_permission_id`` and then defer to the shared
``get_object_permission`` helper so cache entries are shared with the
org / team / key paths."""
from litellm.caching.dual_cache import DualCache
cache = DualCache()
agent_row = MagicMock()
agent_row.object_permission_id = "perm-xyz"
prisma_client = MagicMock()
prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
return_value=agent_row
)
user_api_key_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
agent_id="agent-shared",
)
expected_perm = MagicMock()
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", cache),
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
patch(
"litellm.proxy.auth.auth_checks.get_object_permission",
new_callable=AsyncMock,
return_value=expected_perm,
) as mock_get_perm,
):
result = await MCPRequestHandler._get_agent_object_permission(
user_api_key_auth
)
assert result is expected_perm
mock_get_perm.assert_awaited_once()
assert mock_get_perm.await_args.kwargs["object_permission_id"] == "perm-xyz"
# Second call: the agent_id -> object_permission_id mapping is
# cached, so the agent row is not re-fetched.
prisma_client.db.litellm_agentstable.find_unique.reset_mock()
await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
prisma_client.db.litellm_agentstable.find_unique.assert_not_called()
async def test_get_agent_object_permission_caches_missing_permission(self):
"""When the agent has no ``object_permission_id`` the sentinel must be
cached so subsequent requests do not hit the DB again."""
from litellm.caching.dual_cache import DualCache
cache = DualCache()
agent_row = MagicMock()
agent_row.object_permission_id = None
prisma_client = MagicMock()
prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
return_value=agent_row
)
user_api_key_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
agent_id="agent-no-perm",
)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", cache),
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
patch(
"litellm.proxy.auth.auth_checks.get_object_permission",
new_callable=AsyncMock,
) as mock_get_perm,
):
assert (
await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
is None
)
assert (
await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
is None
)
mock_get_perm.assert_not_awaited()
prisma_client.db.litellm_agentstable.find_unique.assert_awaited_once()
@pytest.mark.asyncio
async def test_tool_permission_servers_included_in_allowed_servers():

14
tests/test_litellm/proxy/auth/test_handle_jwt.py
View file

@ -3341,24 +3341,24 @@ async def test_multi_issuer_jwt_does_not_emit_unscoped_global_warning(
assert JWTHandler._unscoped_jwt_warning_emitted is False
def test_build_decode_kwargs_no_warning_when_issuer_config_scoped(
def test_build_decode_kwargs_warns_for_unscoped_global_fallback_in_mixed_deployment(
monkeypatch, _reset_unscoped_warning_flag, caplog
):
"""When per-issuer config (``LiteLLM_JWTAuth.issuers``) scopes the proxy,
the global path's unscoped-fallback warning must not fire — admins who
intentionally use issuer config without env-var scoping otherwise see a
misleading warning implying their JWT auth is unscoped."""
"""The unscoped-fallback warning must fire even when per-issuer configs
are set. In mixed deployments, tokens whose ``iss`` does not match any
configured issuer fall through to the global path; if env-var scoping is
absent that fallback IS unscoped, and the operator needs to be told."""
import logging
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_ISSUER", raising=False)
caplog.set_level(logging.WARNING)
JWTHandler._build_decode_kwargs(has_issuer_config=True)
JWTHandler._build_decode_kwargs()
matching = [
r
for r in caplog.records
if "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage()
]
assert matching == []
assert len(matching) == 1