feat: litellm shin agent oss staging 05 10 2026 (#27631)
* fix: invalidate cached tag object on tag budget reset (#27481) (#27572) Squash-merged by litellm-agent from oss-agent-shin's PR. * chore(mcp): tighten stdio server registration paths (#27570) Squash-merged by litellm-agent from stuxf's PR. * fix(proxy): clear MCP OpenAPI mappings on server eviction; widen budget cache invalidation Evict OpenAPI tools from global_mcp_tool_registry and strip tool_name_to_mcp_server_name_mapping entries when a server leaves the runtime registry (remove_server and approval-status eviction). Invalidate user_api_key_cache for keys, orgs, and team members on budget-tier spend resets alongside tags. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(mcp): align update_server eviction with remove_server name fallback Document budget-reset test assertion flip (cross-pod cache staleness). Greptile: eviction now pops by server_id then server_name like remove_server; test docstring explains assert_not_awaited -> assert_any_await change. Co-authored-by: Cursor <cursoragent@cursor.com> * Fix org budget cache invalidation --------- Co-authored-by: oss-agent-shin <ext-agent-shin@berri.ai> Co-authored-by: stuxf <70670632+stuxf@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Sameer Kankute
GitHub
parent
6de00e24b4
commit
aa9e7b9808
@ -599,16 +599,57 @@ class MCPServerManager:
|
||||
)
|
||||
raise e
|
||||
|
||||
def _cleanup_server_tool_routing_artifacts(self, server: MCPServer) -> None:
|
||||
"""Drop OpenAPI global tools and name-mapping rows owned by ``server``.
|
||||
|
||||
When a server leaves ``self.registry`` (eviction, ``remove_server``, etc.),
|
||||
OpenAPI tools remain in ``global_mcp_tool_registry`` and
|
||||
``tool_name_to_mcp_server_name_mapping`` unless removed here. Stale
|
||||
mappings make ``_get_mcp_server_from_tool_name`` resolve to a prefix that
|
||||
no longer exists in the live registry.
|
||||
"""
|
||||
from litellm.proxy._experimental.mcp_server.tool_registry import (
|
||||
global_mcp_tool_registry,
|
||||
)
|
||||
|
||||
prefix_root = normalize_server_name(get_server_prefix(server))
|
||||
if server.spec_path and prefix_root:
|
||||
openapi_key_prefix = prefix_root + MCP_TOOL_PREFIX_SEPARATOR
|
||||
global_mcp_tool_registry.unregister_tools_with_prefix(openapi_key_prefix)
|
||||
|
||||
owned_raw: Set[str] = set()
|
||||
for p in iter_known_server_prefixes(server):
|
||||
if p:
|
||||
owned_raw.add(p)
|
||||
if server.name:
|
||||
owned_raw.add(server.name)
|
||||
|
||||
owned_normalized = {normalize_server_name(x) for x in owned_raw}
|
||||
|
||||
stale_mapping_keys: List[str] = []
|
||||
for tool_name, mapped_server in list(
|
||||
self.tool_name_to_mcp_server_name_mapping.items()
|
||||
):
|
||||
if mapped_server in owned_raw:
|
||||
stale_mapping_keys.append(tool_name)
|
||||
elif normalize_server_name(str(mapped_server)) in owned_normalized:
|
||||
stale_mapping_keys.append(tool_name)
|
||||
|
||||
for key in stale_mapping_keys:
|
||||
del self.tool_name_to_mcp_server_name_mapping[key]
|
||||
|
||||
def remove_server(self, mcp_server: LiteLLM_MCPServerTable):
|
||||
"""
|
||||
Remove a server from the registry
|
||||
"""
|
||||
if mcp_server.server_name in self.get_registry():
|
||||
del self.registry[mcp_server.server_name]
|
||||
verbose_logger.debug(f"Removed MCP Server: {mcp_server.server_name}")
|
||||
elif mcp_server.server_id in self.get_registry():
|
||||
del self.registry[mcp_server.server_id]
|
||||
verbose_logger.debug(f"Removed MCP Server: {mcp_server.server_id}")
|
||||
evicted: Optional[MCPServer] = self.registry.pop(mcp_server.server_id, None)
|
||||
if evicted is None and mcp_server.server_name:
|
||||
evicted = self.registry.pop(mcp_server.server_name, None)
|
||||
if evicted is not None:
|
||||
verbose_logger.debug(
|
||||
"Removed MCP Server: %s", mcp_server.server_id or mcp_server.server_name
|
||||
)
|
||||
self._cleanup_server_tool_routing_artifacts(evicted)
|
||||
else:
|
||||
verbose_logger.warning(
|
||||
f"Server ID {mcp_server.server_id} not found in registry"
|
||||
@ -806,6 +847,13 @@ class MCPServerManager:
|
||||
self.initialize_tool_name_to_mcp_server_name_mapping()
|
||||
|
||||
async def add_server(self, mcp_server: LiteLLM_MCPServerTable):
|
||||
# The runtime registry is the allowlist for tool calls and health
|
||||
# probes (which spawn the underlying transport, including stdio
|
||||
# subprocesses). Match the eligibility set used by the bulk DB
|
||||
# filter in reload_servers_from_database() — NULL is legacy and
|
||||
# "approved" is a legacy alias for "active".
|
||||
if mcp_server.approval_status not in (None, "active", "approved"):
|
||||
return
|
||||
try:
|
||||
if mcp_server.server_id not in self.registry:
|
||||
new_server = await self.build_mcp_server_from_table(mcp_server)
|
||||
@ -819,6 +867,16 @@ class MCPServerManager:
|
||||
raise e
|
||||
|
||||
async def update_server(self, mcp_server: LiteLLM_MCPServerTable):
|
||||
# If a previously-active server has been moved out of the active
|
||||
# state, evict any stale registry entry so subsequent tool calls and
|
||||
# health probes can't reach it.
|
||||
if mcp_server.approval_status not in (None, "active", "approved"):
|
||||
evicted = self.registry.pop(mcp_server.server_id, None)
|
||||
if evicted is None and mcp_server.server_name:
|
||||
evicted = self.registry.pop(mcp_server.server_name, None)
|
||||
if evicted is not None:
|
||||
self._cleanup_server_tool_routing_artifacts(evicted)
|
||||
return
|
||||
try:
|
||||
if mcp_server.server_id in self.registry:
|
||||
new_server = await self.build_mcp_server_from_table(mcp_server)
|
||||
|
||||
@ -59,6 +59,22 @@ class MCPToolRegistry:
|
||||
]
|
||||
return list(self.tools.values())
|
||||
|
||||
def unregister_tools_with_prefix(self, prefix: str) -> int:
|
||||
"""Remove tools whose registered name starts with ``prefix``.
|
||||
|
||||
Used when an OpenAPI-backed MCP server leaves the runtime registry so
|
||||
stale tool handlers cannot be invoked after eviction.
|
||||
"""
|
||||
if not prefix:
|
||||
return 0
|
||||
removed = 0
|
||||
for name in list(self.tools.keys()):
|
||||
if name.startswith(prefix):
|
||||
del self.tools[name]
|
||||
removed += 1
|
||||
verbose_logger.debug("Unregistered MCP tool %s", name)
|
||||
return removed
|
||||
|
||||
def convert_tools_to_mcp_sdk_tool_type(
|
||||
self, tools: List[MCPTool]
|
||||
) -> List["MCPToolSDKTool"]:
|
||||
|
||||
@ -87,13 +87,13 @@ class ResetBudgetJob:
|
||||
async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None:
|
||||
"""Drop a stale management-cache entry so the next read fetches from DB.
|
||||
|
||||
Some entity types (notably tags and end-users) are not handled by
|
||||
SpendCounterReseed.from_db, so when a spend counter expires the
|
||||
budget check falls back to ``cached_obj.spend``. If that cached
|
||||
object lingers in ``user_api_key_cache`` past a budget reset, the
|
||||
stale ``.spend`` keeps the entity blocked indefinitely. Deleting
|
||||
the cache entry forces the next auth-time fetch to reload the
|
||||
zeroed row from Postgres.
|
||||
Tags and end-users are not reseeded by ``SpendCounterReseed.from_db``;
|
||||
for those, when the spend counter expires the budget check falls back
|
||||
to ``cached_obj.spend``. Keys, orgs, and team memberships are reseeded
|
||||
from the DB, but auth still may consult ``user_api_key_cache`` objects
|
||||
whose ``.spend`` field can lag a cross-pod DB reset. Deleting the cache
|
||||
entry forces the next auth-time fetch to reload the zeroed row from
|
||||
Postgres.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
@ -113,17 +113,14 @@ class ResetBudgetJob:
|
||||
counter_key_fn: Callable[[Any], str],
|
||||
log_subject: str,
|
||||
extra_where: Optional[dict] = None,
|
||||
cache_key_fn: Optional[Callable[[Any], str]] = None,
|
||||
cache_key_fn: Optional[Callable[[Any], Union[str, List[str]]]] = None,
|
||||
):
|
||||
"""
|
||||
Generic cascade: zero spend on rows whose budget_id is in the reset set.
|
||||
|
||||
``cache_key_fn`` is optional: when provided, after the DB update each
|
||||
matching row's entry in ``user_api_key_cache`` is also dropped. This
|
||||
is required for entities whose spend counter is read with the cached
|
||||
object's ``.spend`` as fallback (tags, end-users) — otherwise the
|
||||
stale cached object pins enforcement to the pre-reset spend until
|
||||
its TTL expires.
|
||||
matching row's entry or entries in ``user_api_key_cache`` are dropped so
|
||||
cached spend cannot stay pinned above the zeroed DB row after a reset.
|
||||
"""
|
||||
budget_ids = [b.budget_id for b in budgets_to_reset if b.budget_id is not None]
|
||||
if not budget_ids:
|
||||
@ -146,7 +143,11 @@ class ResetBudgetJob:
|
||||
for row in rows:
|
||||
await self._invalidate_spend_counter(counter_key_fn(row))
|
||||
if cache_key_fn is not None:
|
||||
await self._invalidate_user_api_key_cache_entry(cache_key_fn(row))
|
||||
cache_keys = cache_key_fn(row)
|
||||
if isinstance(cache_keys, str):
|
||||
cache_keys = [cache_keys]
|
||||
for cache_key in cache_keys:
|
||||
await self._invalidate_user_api_key_cache_entry(cache_key)
|
||||
|
||||
return update_result
|
||||
|
||||
@ -161,6 +162,7 @@ class ResetBudgetJob:
|
||||
table=self.prisma_client.db.litellm_teammembership,
|
||||
counter_key_fn=lambda m: f"spend:team_member:{m.user_id}:{m.team_id}",
|
||||
log_subject="team memberships",
|
||||
cache_key_fn=lambda m: f"{m.team_id}_{m.user_id}",
|
||||
)
|
||||
|
||||
async def reset_budget_for_keys_linked_to_budgets(
|
||||
@ -178,6 +180,7 @@ class ResetBudgetJob:
|
||||
counter_key_fn=lambda k: f"spend:key:{k.token}",
|
||||
log_subject="keys",
|
||||
extra_where={"budget_duration": None, "spend": {"gt": 0}},
|
||||
cache_key_fn=lambda k: k.token,
|
||||
)
|
||||
|
||||
async def reset_budget_for_orgs_linked_to_budgets(
|
||||
@ -192,6 +195,10 @@ class ResetBudgetJob:
|
||||
counter_key_fn=lambda o: f"spend:org:{o.organization_id}",
|
||||
log_subject="orgs",
|
||||
extra_where={"spend": {"gt": 0}},
|
||||
cache_key_fn=lambda o: [
|
||||
f"org_id:{o.organization_id}",
|
||||
f"org_id:{o.organization_id}:with_budget",
|
||||
],
|
||||
)
|
||||
|
||||
async def reset_budget_for_tags_linked_to_budgets(
|
||||
|
||||
@ -142,6 +142,7 @@ if MCP_AVAILABLE:
|
||||
MCPOAuthUserCredentialRequest,
|
||||
MCPOAuthUserCredentialStatus,
|
||||
MCPSubmissionsSummary,
|
||||
MCPTransport,
|
||||
MCPUserCredentialListItem,
|
||||
MCPUserCredentialRequest,
|
||||
MCPUserCredentialResponse,
|
||||
@ -1070,6 +1071,24 @@ if MCP_AVAILABLE:
|
||||
},
|
||||
)
|
||||
|
||||
# stdio servers spawn a local subprocess on the proxy host with the
|
||||
# configured command + args, so accepting them from non-admin callers
|
||||
# would let a team member propose a server config that an admin could
|
||||
# rubber-stamp into local code execution. Restrict stdio submission to
|
||||
# the admin POST /v1/mcp/server path or to config.yaml.
|
||||
if payload.transport == MCPTransport.stdio:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail={
|
||||
"error": (
|
||||
"stdio MCP servers cannot be submitted via the user "
|
||||
"registration workflow. Ask a proxy admin to add this "
|
||||
"server via POST /v1/mcp/server or to declare it in "
|
||||
"config.yaml."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
prisma_client = get_prisma_client_or_throw(
|
||||
"Database not connected. Connect a database to your proxy"
|
||||
)
|
||||
|
||||
@ -1494,6 +1494,7 @@ async def test_add_update_server_with_alias():
|
||||
mock_mcp_server.created_at = None
|
||||
mock_mcp_server.updated_at = None
|
||||
mock_mcp_server.instructions = None
|
||||
mock_mcp_server.approval_status = "active"
|
||||
|
||||
# Add server to manager
|
||||
await test_manager.add_server(mock_mcp_server)
|
||||
@ -1551,6 +1552,7 @@ async def test_add_update_server_without_alias():
|
||||
mock_mcp_server.created_at = None
|
||||
mock_mcp_server.updated_at = None
|
||||
mock_mcp_server.instructions = None
|
||||
mock_mcp_server.approval_status = "active"
|
||||
|
||||
# Add server to manager
|
||||
await test_manager.add_server(mock_mcp_server)
|
||||
@ -1609,6 +1611,7 @@ async def test_add_update_server_fallback_to_server_id():
|
||||
mock_mcp_server.created_at = None
|
||||
mock_mcp_server.updated_at = None
|
||||
mock_mcp_server.instructions = None
|
||||
mock_mcp_server.approval_status = "active"
|
||||
# Add server to manager
|
||||
await test_manager.add_server(mock_mcp_server)
|
||||
|
||||
|
||||
@ -29,7 +29,11 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
_deserialize_json_dict,
|
||||
)
|
||||
from litellm.proxy._types import LiteLLM_MCPServerTable, MCPTransport
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_MCPServerTable,
|
||||
MCPApprovalStatus,
|
||||
MCPTransport,
|
||||
)
|
||||
from litellm.types.mcp import MCPAuth
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer
|
||||
|
||||
@ -3311,5 +3315,126 @@ class TestOAuthDiscoverySSRFGuard:
|
||||
mock_client.get.assert_not_called()
|
||||
|
||||
|
||||
class TestApprovalStatusGate:
|
||||
"""
|
||||
Regression tests for GHSA-gm4g-h72v-jhc3.
|
||||
|
||||
The runtime registry must only contain servers an admin has approved.
|
||||
A non-admin can submit a pending stdio MCP server with an attacker-chosen
|
||||
command/args; before this gate, an admin opening the per-row endpoint
|
||||
triggered ``add_server`` + ``health_check_server``, which spawned the
|
||||
attacker's process under the proxy. The data-layer gate in
|
||||
``add_server`` / ``update_server`` blocks pending and rejected rows
|
||||
from entering the registry regardless of which caller passes them in.
|
||||
"""
|
||||
|
||||
def _make_server(self, server_id: str, approval_status):
|
||||
return LiteLLM_MCPServerTable(
|
||||
server_id=server_id,
|
||||
alias=f"server_{server_id}",
|
||||
description="test",
|
||||
url=None,
|
||||
transport=MCPTransport.stdio,
|
||||
command="python",
|
||||
args=["-c", "print('attacker payload')"],
|
||||
env={},
|
||||
approval_status=approval_status,
|
||||
created_at=datetime.now(),
|
||||
updated_at=datetime.now(),
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"approval_status,expect_in_registry",
|
||||
[
|
||||
(MCPApprovalStatus.pending_review, False),
|
||||
(MCPApprovalStatus.rejected, False),
|
||||
(MCPApprovalStatus.active, True),
|
||||
# Legacy rows: NULL predates the approval workflow; "approved" is
|
||||
# a legacy alias for "active" still present in older deployments.
|
||||
# Both must continue to load to match the DB-level filter in
|
||||
# reload_servers_from_database().
|
||||
(None, True),
|
||||
("approved", True),
|
||||
],
|
||||
)
|
||||
async def test_add_server_respects_approval_status(
|
||||
self, approval_status, expect_in_registry
|
||||
):
|
||||
manager = MCPServerManager()
|
||||
server_id = f"sid-{approval_status}"
|
||||
await manager.add_server(self._make_server(server_id, approval_status))
|
||||
assert (server_id in manager.registry) is expect_in_registry
|
||||
|
||||
async def test_update_server_evicts_when_transitioned_away_from_active(self):
|
||||
# An admin updates a previously-active server to rejected (or pending).
|
||||
# The stale registry entry must be evicted so subsequent tool calls
|
||||
# and health probes can't reach it.
|
||||
manager = MCPServerManager()
|
||||
await manager.add_server(
|
||||
self._make_server("evict-me", MCPApprovalStatus.active)
|
||||
)
|
||||
assert "evict-me" in manager.registry
|
||||
|
||||
await manager.update_server(
|
||||
self._make_server("evict-me", MCPApprovalStatus.rejected)
|
||||
)
|
||||
assert "evict-me" not in manager.registry
|
||||
|
||||
async def test_update_server_eviction_clears_openapi_routing_artifacts(
|
||||
self, tmp_path
|
||||
):
|
||||
"""Rejecting a server must remove its OpenAPI tools and name mappings."""
|
||||
from litellm.proxy._experimental.mcp_server.tool_registry import (
|
||||
global_mcp_tool_registry,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.utils import (
|
||||
add_server_prefix_to_name,
|
||||
get_server_prefix,
|
||||
)
|
||||
|
||||
manager = MCPServerManager()
|
||||
await manager.add_server(
|
||||
self._make_server("evict-openapi", MCPApprovalStatus.active)
|
||||
)
|
||||
assert "evict-openapi" in manager.registry
|
||||
|
||||
server = manager.registry["evict-openapi"]
|
||||
server.spec_path = str(tmp_path / "unused.yaml")
|
||||
prefix = get_server_prefix(server)
|
||||
prefixed = add_server_prefix_to_name("demo_tool", prefix)
|
||||
|
||||
async def _noop_handler(**kwargs):
|
||||
return None
|
||||
|
||||
global_mcp_tool_registry.register_tool(
|
||||
name=prefixed,
|
||||
description="demo",
|
||||
input_schema={"type": "object"},
|
||||
handler=_noop_handler,
|
||||
)
|
||||
manager.tool_name_to_mcp_server_name_mapping["demo_tool"] = prefix
|
||||
manager.tool_name_to_mcp_server_name_mapping[prefixed] = prefix
|
||||
|
||||
await manager.update_server(
|
||||
self._make_server("evict-openapi", MCPApprovalStatus.rejected)
|
||||
)
|
||||
|
||||
assert "evict-openapi" not in manager.registry
|
||||
assert prefixed not in global_mcp_tool_registry.tools
|
||||
assert "demo_tool" not in manager.tool_name_to_mcp_server_name_mapping
|
||||
assert prefixed not in manager.tool_name_to_mcp_server_name_mapping
|
||||
|
||||
async def test_update_server_noop_for_unregistered_pending(self):
|
||||
# update_server called with a pending row that was never registered
|
||||
# should silently return without adding it. Locks in the early-return
|
||||
# so a future refactor can't accidentally route the pending row to
|
||||
# build_mcp_server_from_table.
|
||||
manager = MCPServerManager()
|
||||
await manager.update_server(
|
||||
self._make_server("never-seen", MCPApprovalStatus.pending_review)
|
||||
)
|
||||
assert "never-seen" not in manager.registry
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
|
||||
@ -1527,15 +1527,21 @@ def test_reset_budget_for_tags_linked_to_budgets_invalidates_each_tag_management
|
||||
assert deleted_keys == {"tag:tenant-a", "tag:tenant-b", "tag:tenant-c"}
|
||||
|
||||
|
||||
def test_reset_budget_for_keys_linked_to_budgets_does_not_touch_management_cache(
|
||||
def test_reset_budget_for_keys_linked_to_budgets_invalidates_management_cache(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Cache invalidation is opt-in: keys / orgs / team-members rely on
|
||||
``SpendCounterReseed.from_db`` (which DOES handle their counter keys),
|
||||
so the cache_key_fn hook is intentionally not wired for them. This test
|
||||
locks in that no-op so a future refactor doesn't accidentally start
|
||||
clobbering the key cache (which would cost an extra DB round-trip per
|
||||
reset cycle without fixing anything)."""
|
||||
"""Budget-tier key resets must drop the cached key object (hashed token key).
|
||||
|
||||
Historically this test used ``assert_not_awaited()`` on
|
||||
``user_api_key_cache.async_delete_cache``, reflecting the assumption that
|
||||
``SpendCounterReseed.from_db`` alone kept spend consistent for keys and
|
||||
that invalidating the management cache was unnecessary. That was flipped to
|
||||
``assert_any_await(...)`` because the old invariant fails across pods: a
|
||||
budget reset on one instance can leave another pod's cached key object
|
||||
(including embedded ``.spend``) stale until TTL expiry. Eviction now matches
|
||||
tags/orgs/teams. Do not treat the ``cache_key_fn`` / invalidation wiring as
|
||||
redundant without revisiting that cross-pod consistency story.
|
||||
"""
|
||||
counter_cache = _make_counter_invalidation_job(monkeypatch)
|
||||
|
||||
expired_budget = type("B", (), {"budget_id": "budget-1"})
|
||||
@ -1552,4 +1558,85 @@ def test_reset_budget_for_keys_linked_to_budgets_does_not_touch_management_cache
|
||||
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
|
||||
asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget]))
|
||||
|
||||
counter_cache.user_api_key_cache.async_delete_cache.assert_not_awaited()
|
||||
counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(
|
||||
key="sk-linked"
|
||||
)
|
||||
|
||||
|
||||
def test_reset_budget_for_orgs_linked_to_budgets_invalidates_management_cache(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Org rows use both base and budget-table cache keys — evict both on reset."""
|
||||
counter_cache = _make_counter_invalidation_job(monkeypatch)
|
||||
|
||||
expired_budget = type("B", (), {"budget_id": "budget-1"})
|
||||
linked_org = type("Org", (), {"organization_id": "org-acme"})
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_organizationtable.find_many = AsyncMock(
|
||||
return_value=[linked_org]
|
||||
)
|
||||
prisma_client.db.litellm_organizationtable.update_many = AsyncMock(
|
||||
return_value={"count": 1}
|
||||
)
|
||||
|
||||
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
|
||||
asyncio.run(job.reset_budget_for_orgs_linked_to_budgets([expired_budget]))
|
||||
|
||||
deleted_keys = {
|
||||
call.kwargs.get("key")
|
||||
for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list
|
||||
}
|
||||
assert deleted_keys == {
|
||||
"org_id:org-acme",
|
||||
"org_id:org-acme:with_budget",
|
||||
}
|
||||
|
||||
|
||||
def test_reset_budget_for_team_members_invalidates_management_cache(monkeypatch):
|
||||
"""Team membership cache key matches auth: ``{team_id}_{user_id}``."""
|
||||
counter_cache = _make_counter_invalidation_job(monkeypatch)
|
||||
|
||||
expired_budget = type("B", (), {"budget_id": "budget-1"})
|
||||
membership = type(
|
||||
"Membership",
|
||||
(),
|
||||
{"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"},
|
||||
)
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_teammembership.find_many = AsyncMock(
|
||||
return_value=[membership]
|
||||
)
|
||||
prisma_client.db.litellm_teammembership.update_many = AsyncMock(
|
||||
return_value={"count": 1}
|
||||
)
|
||||
|
||||
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
|
||||
asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget]))
|
||||
|
||||
counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(
|
||||
key="team-x_alice"
|
||||
)
|
||||
|
||||
|
||||
def test_reset_budget_for_tags_linked_to_budgets_management_cache_delete_failure_still_resets(
|
||||
monkeypatch,
|
||||
):
|
||||
"""If ``async_delete_cache`` raises, the DB cascade must still complete."""
|
||||
counter_cache = _make_counter_invalidation_job(monkeypatch)
|
||||
counter_cache.user_api_key_cache.async_delete_cache = AsyncMock(
|
||||
side_effect=RuntimeError("cache unavailable")
|
||||
)
|
||||
|
||||
expired_budget = type("B", (), {"budget_id": "budget-1"})
|
||||
linked_tag = type("Tag", (), {"tag_name": "tenant-42"})
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_tagtable.find_many = AsyncMock(return_value=[linked_tag])
|
||||
prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 1})
|
||||
|
||||
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
|
||||
asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget]))
|
||||
|
||||
prisma_client.db.litellm_tagtable.update_many.assert_awaited_once()
|
||||
|
||||
@ -2491,6 +2491,32 @@ class TestMCPApprovalWorkflow:
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "team" in str(exc_info.value.detail).lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_mcp_server_rejects_stdio_transport(self):
|
||||
# stdio servers spawn a local subprocess on the proxy host. Accepting
|
||||
# them from the non-admin submission endpoint would let a team member
|
||||
# propose a config that an admin could rubber-stamp into local code
|
||||
# execution. Admins use POST /v1/mcp/server or config.yaml instead.
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
|
||||
register_mcp_server,
|
||||
)
|
||||
|
||||
payload = NewMCPServerRequest(
|
||||
alias="local",
|
||||
transport=MCPTransport.stdio,
|
||||
command="python3",
|
||||
args=["-m", "mcp_server_filesystem", "/tmp"],
|
||||
)
|
||||
user_auth = generate_mock_user_api_key_auth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
team_id="team-123",
|
||||
user_id="user-abc",
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await register_mcp_server(payload=payload, user_api_key_dict=user_auth)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "stdio" in str(exc_info.value.detail).lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_mcp_server_sets_pending_review(self):
|
||||
from litellm.proxy._types import MCPApprovalStatus
|
||||
|
||||
Loading…
Reference in New Issue
Block a user