mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
fix(proxy): return resolved object_permission from /key/update
The update response carried object_permission: null even though the row was saved, so the dashboard's key detail view showed zero MCP servers until reload. Echo the resolved object permission the same way /key/info does. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
c8114ba41f
commit
26c1c2c300
2 changed files with 72 additions and 1 deletions
|
|
@ -3385,7 +3385,8 @@ async def update_key_fn(
|
|||
if response is None:
|
||||
raise ValueError("Failed to update key got response = None")
|
||||
|
||||
return {"key": key, **response["data"]}
|
||||
updated_key_info: Final = await attach_object_permission_to_dict({**response["data"]}, prisma_client)
|
||||
return {"key": key, **updated_key_info}
|
||||
# update based on remaining passed in values
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("litellm.proxy.proxy_server.update_key_fn(): Exception occured - %s", e)
|
||||
|
|
|
|||
|
|
@ -808,6 +808,76 @@ async def test_update_key_grandfathers_existing_mcp_servers(monkeypatch):
|
|||
assert result["mcp_servers"] == ["server-a"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_response_echoes_resolved_object_permission(monkeypatch):
|
||||
"""Regression: /key/update must return the resolved object_permission the same
|
||||
way /key/info does. The dashboard spreads the update response over its local
|
||||
key state, so a null object_permission shows "MCP Servers 0" until reload."""
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_ObjectPermissionBase,
|
||||
LitellmUserRoles,
|
||||
UpdateKeyRequest,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
update_key_fn,
|
||||
)
|
||||
|
||||
hashed_token = "c1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
|
||||
existing_key = LiteLLM_VerificationToken(token=hashed_token, user_id="test-user", object_permission_id=None)
|
||||
|
||||
permission_row = MagicMock()
|
||||
permission_row.object_permission_id = "perm-1"
|
||||
permission_row.model_dump.return_value = {
|
||||
"object_permission_id": "perm-1",
|
||||
"mcp_servers": ["server_a"],
|
||||
}
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key)
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=permission_row)
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(return_value=permission_row)
|
||||
mock_prisma_client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma_client.update_data = AsyncMock(
|
||||
return_value={
|
||||
"data": {
|
||||
"token": hashed_token,
|
||||
"user_id": "test-user",
|
||||
"object_permission_id": "perm-1",
|
||||
"object_permission": None,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
_wire_update_key_fn(monkeypatch, existing_key)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
||||
mock_mgr = MagicMock()
|
||||
mock_server = MagicMock()
|
||||
mock_server.server_id = "server_a"
|
||||
mock_mgr.get_registry.return_value = {"server_a": mock_server}
|
||||
mock_mgr.get_allow_all_keys_server_ids.return_value = []
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
mock_mgr,
|
||||
)
|
||||
|
||||
result = await update_key_fn(
|
||||
request=MagicMock(),
|
||||
data=UpdateKeyRequest(
|
||||
key=hashed_token,
|
||||
object_permission=LiteLLM_ObjectPermissionBase(mcp_servers=["server_a"]),
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-admin",
|
||||
user_id="admin-user",
|
||||
),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert result["object_permission"]["mcp_servers"] == ["server_a"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_personal_non_admin_denied_access_groups(
|
||||
monkeypatch,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue