diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 548e3bc3dbf..ed4fb133478 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -536,10 +536,25 @@ class MCPRequestHandler: user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> List[str]: try: - # Get key object permission (already loaded in main auth flow) + # Get key object permission (already loaded in main auth flow, or fetch from DB) key_object_permission = MCPRequestHandler._get_key_object_permission( user_api_key_auth ) + if key_object_permission is None and user_api_key_auth and user_api_key_auth.object_permission_id: + 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 prisma_client is not None: + key_object_permission = await get_object_permission( + object_permission_id=user_api_key_auth.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, + ) if key_object_permission is None: return [] diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index ed8c9a76f60..9c0190bdfec 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3192,7 +3192,11 @@ async def _build_ui_spend_logs_response( if enrich_session_counts: enriched: List[dict] = [] for row in data: - row_dict = row.model_dump() + row_dict = ( + dict(row) + if isinstance(row, dict) + else row.model_dump() + ) sid = row_dict.get("session_id") row_dict["session_total_count"] = count_map.get(sid, 1) if sid else 1 enriched.append(row_dict) diff --git a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py index a3b6a9c6022..c35630176bc 100644 --- a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py @@ -40,10 +40,14 @@ async def test_image_generation_prompt_rerouting(monkeypatch): async def fake_post_call_failure_hook(**_: Any) -> None: return None + async def fake_post_call_success_hook(*, data, user_api_key_dict, response): + return response + fake_proxy_logger = SimpleNamespace( pre_call_hook=fake_pre_call_hook, update_request_status=fake_update_request_status, post_call_failure_hook=fake_post_call_failure_hook, + post_call_success_hook=fake_post_call_success_hook, ) captured_route_request_data: Dict[str, Any] = {} diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index e7da4256182..d65df0087ad 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2996,9 +2996,13 @@ async def test_get_image_non_root_uses_var_lib_assets_dir(monkeypatch): monkeypatch.setenv("LITELLM_NON_ROOT", "true") monkeypatch.delenv("UI_LOGO_PATH", raising=False) - # Mock os.path operations + # Mock os.path operations - exists=False for assets_dir so makedirs gets called + def exists_side_effect(path): + return False if path == "/var/lib/litellm/assets" else True + with patch("litellm.proxy.proxy_server.os.makedirs") as mock_makedirs, \ - patch("litellm.proxy.proxy_server.os.path.exists", return_value=True), \ + patch("litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect), \ + patch("litellm.proxy.proxy_server.os.access", return_value=True), \ patch("litellm.proxy.proxy_server.os.getenv") as mock_getenv, \ patch("litellm.proxy.proxy_server.FileResponse") as mock_file_response: @@ -3038,14 +3042,16 @@ async def test_get_image_non_root_fallback_to_default_logo(monkeypatch): def exists_side_effect(path): exists_calls.append(path) - # Return False for /var/lib/litellm/assets/logo.jpg to trigger fallback - if "/var/lib/litellm/assets/logo.jpg" in path: + # Return False for /var/lib/litellm/assets* so: makedirs is called, logo fallback + # triggers, and we don't return early with cached file + if "/var/lib/litellm/assets" in path: return False return True # Mock os.path operations with patch("litellm.proxy.proxy_server.os.makedirs") as mock_makedirs, \ patch("litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect), \ + patch("litellm.proxy.proxy_server.os.access", return_value=True), \ patch("litellm.proxy.proxy_server.os.getenv") as mock_getenv, \ patch("litellm.proxy.proxy_server.FileResponse") as mock_file_response: