From 82f6d0fe431b1edb65ff6cd1c322fa7c588ba992 Mon Sep 17 00:00:00 2001 From: Alejandro Tapia Date: Thu, 12 Feb 2026 14:47:09 -0800 Subject: [PATCH 01/40] healthcheck-model_id-fix: There was quite a bit of code that needed to be changed since health checks were entirely keyed by model name. This includes, proxy logic, the dashboard, and even networking, because model name was the identifier everywhere. All changed files had tests added to them, which are passing with no regressions. --- litellm/proxy/health_check.py | 12 +- .../health_endpoints/_health_endpoints.py | 8 +- .../litellm_utils_tests/test_health_check.py | 44 ++++ .../proxy/test_health_check_functions.py | 39 ++++ .../ModelsAndEndpointsView.test.tsx | 51 ++++- .../ModelsAndEndpointsView.tsx | 9 +- .../HealthCheckComponent.test.tsx | 147 +++++++++++++ .../model_dashboard/HealthCheckComponent.tsx | 199 ++++++++---------- .../model_dashboard/health_check_columns.tsx | 33 +-- .../src/components/networking.test.ts | 54 +++++ .../src/components/networking.tsx | 10 +- 11 files changed, 469 insertions(+), 137 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.test.tsx diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 427a16a9801..48a20833cb2 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -190,10 +190,15 @@ async def perform_health_check( model: Optional[str] = None, cli_model: Optional[str] = None, details: Optional[bool] = True, + model_id: Optional[str] = None, ): """ Perform a health check on the system. + When model_id is provided, only the deployment with that id is checked + (so models that share the same name but have different ids are checked separately). + When model (name) is provided, all deployments matching that name are checked. + Returns: (bool): True if the health check passes, False otherwise. """ @@ -205,7 +210,12 @@ async def perform_health_check( else: return [], [] - if model is not None: + # Filter by model_id first so a single deployment is checked when id is specified + if model_id is not None: + _by_id = [x for x in model_list if (x.get("model_info") or {}).get("id") == model_id] + if _by_id: + model_list = _by_id + elif model is not None: _new_model_list = [ x for x in model_list if x["litellm_params"]["model"] == model ] diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index da90696ec2d..b844a959622 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -238,7 +238,7 @@ async def health_services_endpoint( # noqa: PLR0915 service_in_success_callbacks = True else: for cb in litellm.success_callback: - if hasattr(cb, 'callback_name') and cb.callback_name == service: + if getattr(cb, "callback_name", None) == service: service_in_success_callbacks = True break cb_id = get_callback_identifier(cb) @@ -732,7 +732,11 @@ async def _perform_health_check_and_save( ): """Helper function to perform health check and save results to database""" healthy_endpoints, unhealthy_endpoints = await perform_health_check( - model_list=model_list, cli_model=cli_model, model=target_model, details=details + model_list=model_list, + cli_model=cli_model, + model=target_model, + details=details, + model_id=model_id, ) # Optionally save health check result to database (non-blocking) diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py index 19882bbe4be..5123b0a7789 100644 --- a/tests/litellm_utils_tests/test_health_check.py +++ b/tests/litellm_utils_tests/test_health_check.py @@ -471,6 +471,50 @@ def test_update_litellm_params_for_health_check(): ) +@pytest.mark.asyncio +async def test_perform_health_check_filters_by_model_id(): + """ + When model_id is passed, only that deployment is checked (not all deployments + that share the same model name). + """ + from litellm.proxy.health_check import perform_health_check + + # Two deployments with same model_name but different ids + model_list = [ + { + "model_name": "gpt-4", + "model_info": {"id": "deployment-id-1"}, + "litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"}, + }, + { + "model_name": "gpt-4", + "model_info": {"id": "deployment-id-2"}, + "litellm_params": {"model": "gpt-4", "api_key": "fake-key-2"}, + }, + ] + + captured_list = [] + + async def mock_perform_health_check(m_list, details=True): + captured_list.append(m_list) + return [{"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]}], [] + + with patch( + "litellm.proxy.health_check._perform_health_check", + side_effect=mock_perform_health_check, + ): + healthy_endpoints, unhealthy_endpoints = await perform_health_check( + model_list=model_list, model_id="deployment-id-2", details=True + ) + + # Only one deployment (deployment-id-2) should have been passed to _perform_health_check + assert len(captured_list) == 1 + assert len(captured_list[0]) == 1 + assert (captured_list[0][0].get("model_info") or {}).get("id") == "deployment-id-2" + assert len(healthy_endpoints) == 1 + assert healthy_endpoints[0]["api_key"] == "fake-key-2" + + @pytest.mark.asyncio async def test_perform_health_check_with_health_check_model(): """ diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index ccae9fb5425..4c91d0ae91e 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -12,6 +12,7 @@ sys.path.insert(0, os.path.abspath("../../..")) from litellm.proxy.health_endpoints._health_endpoints import ( _aggregate_health_check_results, _build_model_param_to_info_mapping, + _perform_health_check_and_save, _save_background_health_checks_to_db, _save_health_check_results_if_changed, _save_health_check_to_db, @@ -466,5 +467,43 @@ async def test_get_all_latest_health_checks_without_model_id(mock_prisma): assert result[0].checked_at == mock_check2.checked_at # Latest +@pytest.mark.asyncio +async def test_perform_health_check_and_save_passes_model_id_to_perform_health_check(): + """Test that _perform_health_check_and_save passes model_id to perform_health_check so health checks run by model id.""" + model_list = [ + { + "model_name": "gpt-4", + "model_info": {"id": "deployment-abc"}, + "litellm_params": {"model": "gpt-4"}, + }, + ] + healthy = [{"model": "gpt-4"}] + unhealthy = [] + + async def mock_perform_health_check(model_list, model=None, cli_model=None, details=True, model_id=None): + return healthy, unhealthy + + with patch( + "litellm.proxy.health_endpoints._health_endpoints.perform_health_check", + side_effect=mock_perform_health_check, + ) as mock_perform: + result = await _perform_health_check_and_save( + model_list=model_list, + target_model=None, + cli_model=None, + details=True, + prisma_client=None, + start_time=0.0, + user_id="user-1", + model_id="deployment-abc", + ) + + mock_perform.assert_called_once() + call_kwargs = mock_perform.call_args[1] + assert call_kwargs["model_id"] == "deployment-abc" + assert result["healthy_count"] == 1 + assert result["unhealthy_count"] == 0 + + if __name__ == "__main__": pytest.main([__file__]) \ No newline at end of file diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx index 1e8eabaea2e..b0df37ad6d8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx @@ -1,6 +1,6 @@ /* @vitest-environment jsdom */ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { render } from "@testing-library/react"; +import { act, render } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import ModelsAndEndpointsView from "./ModelsAndEndpointsView"; @@ -13,6 +13,8 @@ vi.mock("@/components/networking", () => ({ getCallbacksCall: vi.fn().mockResolvedValue({ router_settings: {} }), setCallbacksCall: vi.fn().mockResolvedValue(undefined), getUiSettings: vi.fn().mockResolvedValue({ values: {} }), + latestHealthChecksCall: vi.fn().mockResolvedValue({ latest_health_checks: {} }), + getModelCostMapReloadStatus: vi.fn().mockResolvedValue({}), })); vi.mock("@/app/(dashboard)/models-and-endpoints/components/ModelAnalyticsTab/ModelAnalyticsTab", () => ({ @@ -27,6 +29,14 @@ vi.mock("@/components/add_model/AddModelForm", () => ({ default: () => null, })); +const mockHealthCheckComponent = vi.fn((_props: { all_models_on_proxy?: string[] }) => null); +vi.mock("@/components/model_dashboard/HealthCheckComponent", () => ({ + default: (props: { all_models_on_proxy?: string[] }) => { + mockHealthCheckComponent(props); + return null; + }, +})); + vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ default: () => ({ teams: [], @@ -104,4 +114,43 @@ describe("ModelsAndEndpointsView", () => { ); expect(await findByText("Model Management", {}, { timeout: 10000 })).toBeInTheDocument(); }, 15000); + + it("should pass model IDs (not model names) to HealthCheckComponent as all_models_on_proxy", async () => { + mockHealthCheckComponent.mockClear(); + const modelDataWithIds = { + data: [ + { model_name: "gpt-4", model_info: { id: "deployment-id-1" } }, + { model_name: "gpt-4", model_info: { id: "deployment-id-2" } }, + ], + }; + mockUseModelsInfo.mockReturnValue({ + data: { data: modelDataWithIds.data }, + isLoading: false, + refetch: vi.fn(), + }); + + const queryClient = createQueryClient(); + const { getByRole } = render( + + {}} + premiumUser={false} + teams={[]} + /> + , + ); + + const healthStatusTab = getByRole("tab", { name: "Health Status" }); + await act(async () => { + healthStatusTab.click(); + }); + + expect(mockHealthCheckComponent).toHaveBeenCalled(); + const healthCheckProps = mockHealthCheckComponent.mock.calls[0][0]; + expect(healthCheckProps.all_models_on_proxy).toEqual(["deployment-id-1", "deployment-id-2"]); + expect(healthCheckProps.all_models_on_proxy).not.toContain("gpt-4"); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx index 9d77774cb4c..b697a859dc5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx @@ -98,6 +98,13 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te return modelDataResponse.data.map((model: any) => model.model_name); }, [modelDataResponse?.data]); + const allModelIdsOnProxy = useMemo(() => { + if (!modelDataResponse?.data) return []; + return modelDataResponse.data + .map((model: any) => model.model_info?.id) + .filter((id: string | undefined): id is string => Boolean(id)); + }, [modelDataResponse?.data]); + const getProviderFromModel = (model: string) => { if (modelCostMapData !== null && modelCostMapData !== undefined) { if (typeof modelCostMapData == "object" && model in modelCostMapData) { @@ -397,7 +404,7 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te ({ + individualModelHealthCheckCall: (...args: unknown[]) => mockIndividualModelHealthCheckCall(...args), + latestHealthChecksCall: (...args: unknown[]) => mockLatestHealthChecksCall(...args), +})); + +describe("HealthCheckComponent", () => { + const getDisplayModelName = (model: { model_name?: string }) => model.model_name ?? ""; + + beforeEach(() => { + vi.clearAllMocks(); + mockLatestHealthChecksCall.mockResolvedValue({ latest_health_checks: {} }); + mockIndividualModelHealthCheckCall.mockResolvedValue({ + healthy_count: 1, + unhealthy_count: 0, + healthy_endpoints: [], + unhealthy_endpoints: [], + }); + }); + + it("should render the health check section", async () => { + const modelData = { + data: [ + { + model_name: "gpt-4", + model_info: { id: "deployment-1" }, + litellm_model_name: "gpt-4", + }, + ], + }; + + await act(async () => { + render( + , + ); + }); + await act(async () => { + await new Promise((r) => setTimeout(r, 0)); + }); + + expect(screen.getByText("Model Health Status")).toBeInTheDocument(); + expect( + screen.getByText("Run health checks on individual models to verify they are working correctly"), + ).toBeInTheDocument(); + }); + + it("should call individualModelHealthCheckCall with model id when run health check is triggered", async () => { + const modelData = { + data: [ + { + model_name: "gpt-4", + model_info: { id: "deployment-abc-123" }, + litellm_model_name: "gpt-4", + }, + ], + }; + + render( + , + ); + + const runButtons = screen.getAllByTestId("run-health-check-btn"); + expect(runButtons.length).toBeGreaterThanOrEqual(1); + const runButton = runButtons[0]; + + await act(async () => { + runButton.click(); + }); + + await act(async () => { + await new Promise((r) => setTimeout(r, 50)); + }); + + expect(mockIndividualModelHealthCheckCall).toHaveBeenCalledWith("token-123", "deployment-abc-123"); + expect(mockIndividualModelHealthCheckCall).not.toHaveBeenCalledWith("token-123", "gpt-4"); + }); + + it("should key health status by model id and show status from latest_health_checks by model_id", async () => { + const modelData = { + data: [ + { + model_name: "gpt-4", + model_info: { id: "id-alpha" }, + litellm_model_name: "gpt-4", + }, + { + model_name: "gpt-4", + model_info: { id: "id-beta" }, + litellm_model_name: "gpt-4", + }, + ], + }; + + mockLatestHealthChecksCall.mockResolvedValue({ + latest_health_checks: { + "id-alpha": { + status: "healthy", + checked_at: "2024-01-15T10:00:00Z", + error_message: null, + }, + "id-beta": { + status: "unhealthy", + checked_at: "2024-01-15T10:05:00Z", + error_message: "Connection failed", + }, + }, + }); + + await act(async () => { + render( + , + ); + }); + + await act(async () => { + await new Promise((r) => setTimeout(r, 0)); + }); + + expect(mockLatestHealthChecksCall).toHaveBeenCalledWith("token"); + const healthyBadges = screen.getAllByText("healthy"); + const unhealthyBadges = screen.getAllByText("unhealthy"); + expect(healthyBadges.length).toBeGreaterThanOrEqual(1); + expect(unhealthyBadges.length).toBeGreaterThanOrEqual(1); + }); +}); diff --git a/ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx b/ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx index 6fa83494408..b4bb1019dd6 100644 --- a/ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx +++ b/ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx @@ -53,31 +53,33 @@ const HealthCheckComponent: React.FC = ({ const healthTableRef = useRef>(null); - // Initialize health statuses on component mount + // Initialize health statuses on component mount (keyed by model id) useEffect(() => { if (!accessToken || !modelData?.data) return; const initializeHealthStatuses = async () => { const healthStatusMap: { [key: string]: HealthStatus } = {}; - // Initialize all models with default state using model names + // Initialize all models with default state using model ids modelData.data.forEach((model: any) => { - const modelName = model.model_name; - healthStatusMap[modelName] = { - status: "none", - lastCheck: "None", - lastSuccess: "None", - loading: false, - error: undefined, - fullError: undefined, - successResponse: undefined, - }; + const modelId = model.model_info?.id; + if (modelId) { + healthStatusMap[modelId] = { + status: "none", + lastCheck: "None", + lastSuccess: "None", + loading: false, + error: undefined, + fullError: undefined, + successResponse: undefined, + }; + } }); try { const latestHealthChecks = await latestHealthChecksCall(accessToken); - // Override with actual database data if it exists + // Override with actual database data if it exists (latest_health_checks is keyed by model_id) if ( latestHealthChecks && latestHealthChecks.latest_health_checks && @@ -86,32 +88,28 @@ const HealthCheckComponent: React.FC = ({ Object.entries(latestHealthChecks.latest_health_checks).forEach(([key, checkData]: [string, any]) => { if (!checkData) return; - let targetModelName: string | null = null; + let targetModelId: string | null = null; - // The key could be either model_id or model_name, try both approaches - const directModelMatch = modelData.data.find((m: any) => m.model_name === key); - if (directModelMatch) { - targetModelName = directModelMatch.model_name; + // The key is model_id from the backend; fallback to matching by model_name for legacy data + const modelByIdMatch = modelData.data.find((m: any) => m.model_info && m.model_info.id === key); + if (modelByIdMatch) { + targetModelId = modelByIdMatch.model_info.id; } else { - // If not a direct match, treat as model_id and find the corresponding model - const modelByIdMatch = modelData.data.find((m: any) => m.model_info && m.model_info.id === key); - if (modelByIdMatch) { - targetModelName = modelByIdMatch.model_name; - } else { - // Check if checkData contains model_name and use that - if (checkData.model_name) { - const modelByNameInData = modelData.data.find((m: any) => m.model_name === checkData.model_name); - if (modelByNameInData) { - targetModelName = modelByNameInData.model_name; - } + const directModelMatch = modelData.data.find((m: any) => m.model_name === key); + if (directModelMatch?.model_info?.id) { + targetModelId = directModelMatch.model_info.id; + } else if (checkData.model_name) { + const modelByNameInData = modelData.data.find((m: any) => m.model_name === checkData.model_name); + if (modelByNameInData?.model_info?.id) { + targetModelId = modelByNameInData.model_info.id; } } } - if (targetModelName) { + if (targetModelId) { const fullError = checkData.error_message || undefined; - healthStatusMap[targetModelName] = { + healthStatusMap[targetModelId] = { status: checkData.status || "unknown", lastCheck: checkData.checked_at ? new Date(checkData.checked_at).toLocaleString() : "None", lastSuccess: @@ -246,33 +244,31 @@ const HealthCheckComponent: React.FC = ({ return cleaned.length > 100 ? cleaned.substring(0, 97) + "..." : cleaned; }; - const runIndividualHealthCheck = async (modelName: string) => { + const runIndividualHealthCheck = async (modelId: string) => { if (!accessToken) return; setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { - ...prev[modelName], + [modelId]: { + ...prev[modelId], loading: true, status: "checking", }, })); try { - // Run the health check and process the response directly - const response = await individualModelHealthCheckCall(accessToken, modelName); + const response = await individualModelHealthCheckCall(accessToken, modelId); const currentTime = new Date().toLocaleString(); - // Check if there are any unhealthy endpoints (which means this specific model failed) if (response.unhealthy_count > 0 && response.unhealthy_endpoints && response.unhealthy_endpoints.length > 0) { const rawError = response.unhealthy_endpoints[0]?.error || "Health check failed"; const errorMessage = extractMeaningfulError(rawError); setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { + [modelId]: { status: "unhealthy", lastCheck: currentTime, - lastSuccess: prev[modelName]?.lastSuccess || "None", + lastSuccess: prev[modelId]?.lastSuccess || "None", loading: false, error: errorMessage, fullError: rawError, @@ -281,7 +277,7 @@ const HealthCheckComponent: React.FC = ({ } else { setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { + [modelId]: { status: "healthy", lastCheck: currentTime, lastSuccess: currentTime, @@ -291,41 +287,33 @@ const HealthCheckComponent: React.FC = ({ })); } - // Refresh health status from database to get the saved check data including timestamp try { const latestHealthChecks = await latestHealthChecksCall(accessToken); + const checkData = latestHealthChecks.latest_health_checks?.[modelId]; - // Find the model ID for this model name to look up database data - const model = modelData.data.find((m: any) => m.model_name === modelName); - if (model) { - const modelId = model.model_info.id; - const checkData = latestHealthChecks.latest_health_checks?.[modelId]; - - if (checkData) { - const fullError = checkData.error_message || undefined; - setModelHealthStatuses((prev) => ({ - ...prev, - [modelName]: { - status: checkData.status || prev[modelName]?.status || "unknown", - lastCheck: checkData.checked_at - ? new Date(checkData.checked_at).toLocaleString() - : prev[modelName]?.lastCheck || "None", - lastSuccess: - checkData.status === "healthy" - ? checkData.checked_at - ? new Date(checkData.checked_at).toLocaleString() - : prev[modelName]?.lastSuccess || "None" - : prev[modelName]?.lastSuccess || "None", - loading: false, - error: fullError ? extractMeaningfulError(fullError) : prev[modelName]?.error, - fullError: fullError || prev[modelName]?.fullError, - successResponse: checkData.status === "healthy" ? checkData : prev[modelName]?.successResponse, - }, - })); - } + if (checkData) { + const fullError = checkData.error_message || undefined; + setModelHealthStatuses((prev) => ({ + ...prev, + [modelId]: { + status: checkData.status || prev[modelId]?.status || "unknown", + lastCheck: checkData.checked_at + ? new Date(checkData.checked_at).toLocaleString() + : prev[modelId]?.lastCheck || "None", + lastSuccess: + checkData.status === "healthy" + ? checkData.checked_at + ? new Date(checkData.checked_at).toLocaleString() + : prev[modelId]?.lastSuccess || "None" + : prev[modelId]?.lastSuccess || "None", + loading: false, + error: fullError ? extractMeaningfulError(fullError) : prev[modelId]?.error, + fullError: fullError || prev[modelId]?.fullError, + successResponse: checkData.status === "healthy" ? checkData : prev[modelId]?.successResponse, + }, + })); } } catch (dbError) { - // Ignore database errors - we already have the health check result from the API call console.debug("Could not fetch updated status from database (non-critical):", dbError); } } catch (error) { @@ -334,10 +322,10 @@ const HealthCheckComponent: React.FC = ({ const errorMessage = extractMeaningfulError(rawError); setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { + [modelId]: { status: "unhealthy", lastCheck: currentTime, - lastSuccess: prev[modelName]?.lastSuccess || "None", + lastSuccess: prev[modelId]?.lastSuccess || "None", loading: false, error: errorMessage, fullError: rawError, @@ -349,11 +337,10 @@ const HealthCheckComponent: React.FC = ({ const runAllHealthChecks = async () => { const modelsToCheck = selectedModelsForHealth.length > 0 ? selectedModelsForHealth : all_models_on_proxy; - // Set all models to loading state const loadingStatuses = modelsToCheck.reduce( - (acc, modelName) => { - acc[modelName] = { - ...modelHealthStatuses[modelName], + (acc, modelId) => { + acc[modelId] = { + ...modelHealthStatuses[modelId], loading: true, status: "checking", }; @@ -364,30 +351,25 @@ const HealthCheckComponent: React.FC = ({ setModelHealthStatuses((prev) => ({ ...prev, ...loadingStatuses })); - // Store results from individual health checks const healthCheckResults: { [key: string]: any } = {}; - // Run all health checks in parallel and collect results - const healthCheckPromises = modelsToCheck.map(async (modelName) => { + const healthCheckPromises = modelsToCheck.map(async (modelId) => { if (!accessToken) return; try { - // Run the health check and store the result - const response = await individualModelHealthCheckCall(accessToken, modelName); - healthCheckResults[modelName] = response; + const response = await individualModelHealthCheckCall(accessToken, modelId); + healthCheckResults[modelId] = response; - // Update status immediately based on response const currentTime = new Date().toLocaleString(); - // Check if there are any unhealthy endpoints (which means this specific model failed) if (response.unhealthy_count > 0 && response.unhealthy_endpoints && response.unhealthy_endpoints.length > 0) { const rawError = response.unhealthy_endpoints[0]?.error || "Health check failed"; const errorMessage = extractMeaningfulError(rawError); setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { + [modelId]: { status: "unhealthy", lastCheck: currentTime, - lastSuccess: prev[modelName]?.lastSuccess || "None", + lastSuccess: prev[modelId]?.lastSuccess || "None", loading: false, error: errorMessage, fullError: rawError, @@ -396,7 +378,7 @@ const HealthCheckComponent: React.FC = ({ } else { setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { + [modelId]: { status: "healthy", lastCheck: currentTime, lastSuccess: currentTime, @@ -406,17 +388,16 @@ const HealthCheckComponent: React.FC = ({ })); } } catch (error) { - console.error(`Health check failed for ${modelName}:`, error); - // Set error status for failed health checks + console.error(`Health check failed for model id ${modelId}:`, error); const currentTime = new Date().toLocaleString(); const rawError = error instanceof Error ? error.message : String(error); const errorMessage = extractMeaningfulError(rawError); setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { + [modelId]: { status: "unhealthy", lastCheck: currentTime, - lastSuccess: prev[modelName]?.lastSuccess || "None", + lastSuccess: prev[modelId]?.lastSuccess || "None", loading: false, error: errorMessage, fullError: rawError, @@ -425,27 +406,21 @@ const HealthCheckComponent: React.FC = ({ } }); - // Wait for all health checks to complete await Promise.allSettled(healthCheckPromises); - // Refresh health statuses from database to get the saved check data including timestamps try { if (!accessToken) return; const latestHealthChecks = await latestHealthChecksCall(accessToken); if (latestHealthChecks.latest_health_checks) { - // Update health statuses from database, which should have the most accurate saved data Object.entries(latestHealthChecks.latest_health_checks).forEach(([modelId, checkData]: [string, any]) => { - // Find the model name for this model ID - const model = modelData.data.find((m: any) => m.model_info.id === modelId); - if (model && modelsToCheck.includes(model.model_name) && checkData) { - const modelName = model.model_name; + if (modelsToCheck.includes(modelId) && checkData) { const fullError = checkData.error_message || undefined; setModelHealthStatuses((prev) => { - const currentStatus = prev[modelName]; + const currentStatus = prev[modelId]; return { ...prev, - [modelName]: { + [modelId]: { status: checkData.status || currentStatus?.status || "unknown", lastCheck: checkData.checked_at ? new Date(checkData.checked_at).toLocaleString() @@ -468,15 +443,14 @@ const HealthCheckComponent: React.FC = ({ } } catch (dbError) { console.warn("Failed to fetch updated health statuses from database (non-critical):", dbError); - // This is non-critical - we already have the health check results from the API calls } }; - const handleModelSelection = (modelName: string, checked: boolean) => { + const handleModelSelection = (modelId: string, checked: boolean) => { if (checked) { - setSelectedModelsForHealth((prev) => [...prev, modelName]); + setSelectedModelsForHealth((prev) => [...prev, modelId]); } else { - setSelectedModelsForHealth((prev) => prev.filter((name) => name !== modelName)); + setSelectedModelsForHealth((prev) => prev.filter((id) => id !== modelId)); setAllModelsSelected(false); } }; @@ -580,8 +554,9 @@ const HealthCheckComponent: React.FC = ({ teams, )} data={modelData.data.map((model: any) => { - const modelName = model.model_name; - const healthStatus = modelHealthStatuses[modelName] || { + const modelId = model.model_info?.id; + const healthStatus = modelId ? modelHealthStatuses[modelId] : null; + const status = healthStatus || { status: "none", lastCheck: "None", loading: false, @@ -591,12 +566,12 @@ const HealthCheckComponent: React.FC = ({ model_info: model.model_info, provider: model.provider, litellm_model_name: model.litellm_model_name, - health_status: healthStatus.status, - last_check: healthStatus.lastCheck, - last_success: healthStatus.lastSuccess || "None", - health_loading: healthStatus.loading, - health_error: healthStatus.error, - health_full_error: healthStatus.fullError, + health_status: status.status, + last_check: status.lastCheck, + last_success: status.lastSuccess || "None", + health_loading: status.loading, + health_error: status.error, + health_full_error: status.fullError, }; })} isLoading={false} diff --git a/ui/litellm-dashboard/src/components/model_dashboard/health_check_columns.tsx b/ui/litellm-dashboard/src/components/model_dashboard/health_check_columns.tsx index 3e8ae662ad4..396afb7ed0b 100644 --- a/ui/litellm-dashboard/src/components/model_dashboard/health_check_columns.tsx +++ b/ui/litellm-dashboard/src/components/model_dashboard/health_check_columns.tsx @@ -35,9 +35,9 @@ export const healthCheckColumns = ( modelHealthStatuses: { [key: string]: HealthStatus }, selectedModelsForHealth: string[], allModelsSelected: boolean, - handleModelSelection: (modelName: string, checked: boolean) => void, + handleModelSelection: (modelId: string, checked: boolean) => void, handleSelectAll: (checked: boolean) => void, - runIndividualHealthCheck: (modelName: string) => void, + runIndividualHealthCheck: (modelId: string) => void, getStatusBadge: (status: string) => JSX.Element, getDisplayModelName: (model: any) => string, showErrorModal?: (modelName: string, cleanedError: string, fullError: string) => void, @@ -62,14 +62,14 @@ export const healthCheckColumns = ( sortingFn: "alphanumeric", cell: ({ row }) => { const model = row.original; - const modelName = model.model_name; - const isSelected = selectedModelsForHealth.includes(modelName); + const modelId = model.model_info?.id ?? ""; + const isSelected = selectedModelsForHealth.includes(modelId); return (
handleModelSelection(modelName, e.target.checked)} + onChange={(e) => handleModelSelection(modelId, e.target.checked)} onClick={(e) => e.stopPropagation()} /> @@ -169,8 +169,9 @@ export const healthCheckColumns = ( ); } - const modelName = model.model_name; - const hasSuccessResponse = healthStatus.status === "healthy" && modelHealthStatuses[modelName]?.successResponse; + const modelId = model.model_info?.id ?? ""; + const displayName = getDisplayModelName(model) || model.model_name; + const hasSuccessResponse = healthStatus.status === "healthy" && modelHealthStatuses[modelId]?.successResponse; return (
@@ -178,7 +179,7 @@ export const healthCheckColumns = ( {hasSuccessResponse && showSuccessModal && ( +
+ )} + +
+ +
+ + +
+ + {data.keyName} + + + Key ID: {data.keyId} + +
+ {canModifyKey && ( + + + + + + + + + )} + + + + + } /> + } + truncate + copyable + defaultUserIdCheck + /> + + + + + + } /> + } + truncate + copyable + defaultUserIdCheck + /> + + + + + + } /> + } /> + + +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 94ca90b9630..378f5b3872b 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -1,13 +1,13 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import useTeams from "@/app/(dashboard)/hooks/useTeams"; -import { formatNumberWithCommas, copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; import { mapEmptyStringToNull } from "@/utils/keyUpdateUtils"; -import { ArrowLeftIcon, RefreshIcon, TrashIcon } from "@heroicons/react/outline"; +import { ArrowLeftIcon } from "@heroicons/react/outline"; import { Badge, Button, Card, Grid, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react"; -import { Button as AntdButton, Form, Tag, Tooltip } from "antd"; -import { CheckIcon, CopyIcon } from "lucide-react"; +import { Form, Tag } from "antd"; +import { KeyInfoHeader } from "./KeyInfoHeader"; import { useEffect, useState } from "react"; -import { isProxyAdminRole, isUserTeamAdminForSingleTeam, rolesWithWriteAccess } from "../../utils/roles"; +import { isProxyAdminRole, isUserTeamAdminForSingleTeam } from "../../utils/roles"; import { mapDisplayToInternalNames, mapInternalToDisplayNames } from "../callback_info_helpers"; import AutoRotationView from "../common_components/AutoRotationView"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; @@ -54,8 +54,6 @@ export default function KeyInfoView({ const [deleteLoading, setDeleteLoading] = useState(false); const [deleteConfirmInput, setDeleteConfirmInput] = useState(""); const [isRegenerateModalOpen, setIsRegenerateModalOpen] = useState(false); - const [copiedStates, setCopiedStates] = useState>({}); - // Add local state to maintain key data and track regeneration const [currentKeyData, setCurrentKeyData] = useState(keyData); const [lastRegeneratedAt, setLastRegeneratedAt] = useState(null); @@ -284,16 +282,6 @@ export default function KeyInfoView({ } }; - const copyToClipboard = async (text: string, key: string) => { - const success = await utilCopyToClipboard(text); - if (success) { - setCopiedStates((prev) => ({ ...prev, [key]: true })); - setTimeout(() => { - setCopiedStates((prev) => ({ ...prev, [key]: false })); - }, 2000); - } - }; - const handleRegenerateKeyUpdate = (updatedKeyData: Partial) => { // Update local state immediately with ALL the new data setCurrentKeyData((prevData) => { @@ -346,79 +334,29 @@ export default function KeyInfoView({ return (
-
-
- - {currentKeyData.key_alias || "Virtual Key"} - -
-
- Key ID - {currentKeyData.token_id || currentKeyData.token} -
- : } - onClick={() => copyToClipboard(currentKeyData.token_id || currentKeyData.token, "key-id")} - className={`ml-2 transition-all duration-200${copiedStates["key-id"] - ? "text-green-600 bg-green-50 border-green-200" - : "text-gray-500 hover:text-gray-700 hover:bg-gray-100" - }`} - /> -
- - {/* Add timestamp and regeneration indicator */} -
- - {currentKeyData.updated_at && currentKeyData.updated_at !== currentKeyData.created_at - ? `Updated: ${formatTimestamp(currentKeyData.updated_at)}` - : `Created: ${formatTimestamp(currentKeyData.created_at)}`} - - - {isRecentlyRegenerated && ( - - Recently Regenerated - - )} - - {lastRegeneratedAt && ( - - Regenerated - - )} -
-
- {canModifyKey && ( -
- - - - - - -
- )} -
+ setIsRegenerateModalOpen(true)} + onDelete={() => setIsDeleteModalOpen(true)} + canModifyKey={canModifyKey} + backButtonText={backButtonText} + regenerateDisabled={!premiumUser} + regenerateTooltip={ + !premiumUser + ? "This is a LiteLLM Enterprise feature, and requires a valid key to use." + : undefined + } + /> {/* Add RegenerateKeyModal */} Date: Tue, 24 Feb 2026 17:14:37 -0800 Subject: [PATCH 07/40] [Fix] Enrich failure spend logs with key/team metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Failure spend logs were missing key metadata (key alias, user ID, team ID, team alias) in two scenarios: 1. Auth errors (401 ProxyException): auth_exception_handler creates a minimal UserAPIKeyAuth with only api_key and request_route set — all other fields are null. The failure hook now looks up the full key object from cache/DB using the key hash to populate the missing fields. 2. Post-auth failures (provider errors, rate limits): key fields are present but team_alias is always null because LiteLLM_VerificationTokenView SQL view does not include team_alias. The failure hook now looks up the team object from cache to populate team_alias. Both lookups are non-fatal and wrapped in try/except. Co-Authored-By: Claude Sonnet 4.6 --- .../proxy/hooks/proxy_track_cost_callback.py | 73 +++++- .../hooks/test_proxy_track_cost_callback.py | 217 ++++++++++++++++++ 2 files changed, 289 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 815d64f22ad..80e202e9578 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -12,7 +12,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.auth.auth_checks import log_db_metrics +from litellm.proxy.auth.auth_checks import get_key_object, get_team_object, log_db_metrics from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( @@ -76,6 +76,10 @@ class _ProxyDBLogger(CustomLogger): traceback_str=traceback_str, ) + _metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info( + metadata=_metadata, + ) + existing_metadata: dict = request_data.get("metadata", None) or {} existing_metadata.update(_metadata) @@ -255,6 +259,73 @@ class _ProxyDBLogger(CustomLogger): "Error in tracking cost callback - %s", str(e) ) + @staticmethod + async def _enrich_failure_metadata_with_key_info(metadata: dict) -> dict: + """ + Enriches failure spend log metadata by looking up the key object (and team object) + from cache/DB when key fields are missing. + + This handles two scenarios: + 1. Auth errors (401): UserAPIKeyAuth is created with only api_key set, all other + fields are null. We look up the full key object to fill in alias, user_id, + team_id, etc. + 2. Post-auth failures (provider errors, rate limits): key fields are populated + but team_alias is missing because LiteLLM_VerificationTokenView SQL view + doesn't include it. We look up the team object to fill in team_alias. + """ + api_key_hash = metadata.get("user_api_key") + if not api_key_hash: + return metadata + + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + # Step 1: If key fields are missing, look up the full key object + if metadata.get("user_api_key_alias") is None: + try: + key_obj = await get_key_object( + hashed_token=api_key_hash, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if metadata.get("user_api_key_alias") is None: + metadata["user_api_key_alias"] = key_obj.key_alias + if metadata.get("user_api_key_user_id") is None: + metadata["user_api_key_user_id"] = key_obj.user_id + if metadata.get("user_api_key_team_id") is None: + metadata["user_api_key_team_id"] = key_obj.team_id + if metadata.get("user_api_key_org_id") is None: + metadata["user_api_key_org_id"] = key_obj.org_id + except Exception: + verbose_proxy_logger.debug( + "Failed to enrich failure metadata with key info for api_key=%s", + api_key_hash, + ) + + # Step 2: If team_id is known but team_alias is missing, look up the team object + team_id = metadata.get("user_api_key_team_id") + if team_id and metadata.get("user_api_key_team_alias") is None: + try: + team_obj = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_cache_only=True, + ) + if team_obj.team_alias is not None: + metadata["user_api_key_team_alias"] = team_obj.team_alias + except Exception: + verbose_proxy_logger.debug( + "Failed to enrich failure metadata with team_alias for team_id=%s", + team_id, + ) + return metadata + @staticmethod def _should_track_errors_in_db(): """ diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index e8765cf78ca..c46b8df5efc 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -169,6 +169,223 @@ async def test_track_cost_callback_skips_when_no_standard_logging_object(): mock_proxy_logging.failed_tracking_alert.assert_not_called() +@pytest.mark.asyncio +async def test_enrich_failure_metadata_with_team_alias(): + """ + When team_id is set but team_alias is missing (and key_alias is present), + _enrich_failure_metadata_with_key_info should look up the team from cache + and populate user_api_key_team_alias. + """ + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "my-team-alias" + + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ): + metadata = { + "user_api_key": "hashed_key", + "user_api_key_alias": "my-key-alias", # already set + "user_api_key_team_id": "test_team_id", + "user_api_key_team_alias": None, + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) + assert result["user_api_key_team_alias"] == "my-team-alias" + + +@pytest.mark.asyncio +async def test_enrich_failure_metadata_with_full_key_lookup(): + """ + When all key fields are null (auth error 401 scenario), _enrich_failure_metadata_with_key_info + should look up the key object from cache/DB and populate alias, user_id, team_id, + then look up the team to get team_alias. + """ + mock_key_obj = MagicMock() + mock_key_obj.key_alias = "fetched-key-alias" + mock_key_obj.user_id = "fetched-user-id" + mock_key_obj.team_id = "fetched-team-id" + mock_key_obj.org_id = "fetched-org-id" + + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "fetched-team-alias" + + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + return_value=mock_key_obj, + ), patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ): + metadata = { + "user_api_key": "hashed_key", + "user_api_key_alias": None, # all null - simulates auth error path + "user_api_key_user_id": None, + "user_api_key_team_id": None, + "user_api_key_team_alias": None, + "user_api_key_org_id": None, + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) + assert result["user_api_key_alias"] == "fetched-key-alias" + assert result["user_api_key_user_id"] == "fetched-user-id" + assert result["user_api_key_team_id"] == "fetched-team-id" + assert result["user_api_key_org_id"] == "fetched-org-id" + assert result["user_api_key_team_alias"] == "fetched-team-alias" + + +@pytest.mark.asyncio +async def test_enrich_failure_metadata_skips_when_team_alias_present(): + """ + When team_alias is already populated, _enrich_failure_metadata_with_key_info + should not perform a team cache lookup. + """ + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + ) as mock_get_key, patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + metadata = { + "user_api_key": "hashed_key", + "user_api_key_alias": "existing-alias", + "user_api_key_team_id": "test_team_id", + "user_api_key_team_alias": "already-set", + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) + assert result["user_api_key_team_alias"] == "already-set" + mock_get_key.assert_not_called() + mock_get_team.assert_not_called() + + +@pytest.mark.asyncio +async def test_enrich_failure_metadata_skips_when_no_api_key(): + """ + When api_key hash is absent, _enrich_failure_metadata_with_key_info should + not perform any lookups. + """ + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + ) as mock_get_key: + metadata = { + "user_api_key": None, + "user_api_key_alias": None, + "user_api_key_team_id": None, + "user_api_key_team_alias": None, + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) + mock_get_key.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_enriches_auth_error_metadata(): + """ + Simulates a 401 ProxyException (e.g. can_key_call_model). In this case + UserAPIKeyAuth is created with only api_key set. The failure hook should + look up the key and team from cache/DB to populate all missing fields. + """ + logger = _ProxyDBLogger() + + # This is what auth_exception_handler creates for 401 errors + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed_key", + # key_alias, user_id, team_id, team_alias are all None + ) + + request_data = { + "model": "claude-haiku-4-5", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "litellm_params": {}, + } + + mock_key_obj = MagicMock() + mock_key_obj.key_alias = "my-key-alias" + mock_key_obj.user_id = "my-user-id" + mock_key_obj.team_id = "my-team-id" + mock_key_obj.org_id = None + + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "my-team-alias" + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + return_value=mock_key_obj, + ), patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ): + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("401 - model not allowed"), + user_api_key_dict=user_api_key_dict, + ) + + mock_update_database.assert_called_once() + call_args = mock_update_database.call_args[1] + metadata = call_args["kwargs"]["litellm_params"]["metadata"] + assert metadata["user_api_key_alias"] == "my-key-alias" + assert metadata["user_api_key_user_id"] == "my-user-id" + assert metadata["user_api_key_team_id"] == "my-team-id" + assert metadata["user_api_key_team_alias"] == "my-team-alias" + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_enriches_missing_team_alias(): + """ + When user_api_key_dict has a team_id but no team_alias, async_post_call_failure_hook + should look up the team from cache and populate user_api_key_team_alias in the + spend log metadata written to the DB. + """ + logger = _ProxyDBLogger() + + user_api_key_dict = UserAPIKeyAuth( + api_key="test_api_key", + key_alias="test_alias", + user_id="test_user_id", + team_id="test_team_id", + team_alias=None, # Missing - simulates regular key auth where SQL view omits team_alias + ) + + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "litellm_params": {}, + } + + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "enriched-team-alias" + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ): + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("Provider rate limit"), + user_api_key_dict=user_api_key_dict, + ) + + mock_update_database.assert_called_once() + call_args = mock_update_database.call_args[1] + metadata = call_args["kwargs"]["litellm_params"]["metadata"] + assert metadata["user_api_key_team_alias"] == "enriched-team-alias" + assert metadata["user_api_key_team_id"] == "test_team_id" + + @pytest.mark.asyncio @pytest.mark.parametrize("model_value", [None, ""]) async def test_track_cost_callback_skips_for_falsy_model_and_no_slo(model_value): From b1423bc8bf462388f7920a06244358bc6f03d4b2 Mon Sep 17 00:00:00 2001 From: Alejandro Tapia Date: Tue, 24 Feb 2026 17:37:00 -0800 Subject: [PATCH 08/40] removed extra comma typo --- litellm/proxy/health_endpoints/_health_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index f865e67c7a9..f3bed3656f6 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -746,7 +746,7 @@ async def _perform_health_check_and_save( cli_model=cli_model, model=target_model, details=details, - max_concurrency=max_concurrency,, + max_concurrency=max_concurrency, model_id=model_id, ) From 4e84e4c60729bdd63d03538ed6143e84a2bffe32 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 24 Feb 2026 18:16:04 -0800 Subject: [PATCH 09/40] fix(ui): show real tool names in logs for Anthropic format tools (#22048) --- .../view_logs/ToolsSection/utils.ts | 42 ++++++++++++++----- 1 file changed, 32 insertions(+), 10 deletions(-) diff --git a/ui/litellm-dashboard/src/components/view_logs/ToolsSection/utils.ts b/ui/litellm-dashboard/src/components/view_logs/ToolsSection/utils.ts index 33b21297c43..351bbf51169 100644 --- a/ui/litellm-dashboard/src/components/view_logs/ToolsSection/utils.ts +++ b/ui/litellm-dashboard/src/components/view_logs/ToolsSection/utils.ts @@ -47,9 +47,9 @@ function extractToolsFromRequest(log: LogEntry): ToolDefinition[] { */ function extractToolCallsFromResponse(log: LogEntry): ToolCall[] { const responseData = parseData(log.response); - + if (!responseData || typeof responseData !== "object") return []; - + // OpenAI format: response.choices[0].message.tool_calls const choices = responseData.choices; if (Array.isArray(choices) && choices.length > 0) { @@ -59,7 +59,24 @@ function extractToolCallsFromResponse(log: LogEntry): ToolCall[] { return message.tool_calls; } } - + + // Anthropic format: response.content[].type === "tool_use" + if (Array.isArray(responseData.content)) { + const toolUseBlocks = responseData.content.filter( + (block: any) => block.type === "tool_use" + ); + if (toolUseBlocks.length > 0) { + return toolUseBlocks.map((block: any) => ({ + id: block.id, + type: "function", + function: { + name: block.name, + arguments: JSON.stringify(block.input || {}), + }, + })); + } + } + return []; } @@ -106,15 +123,20 @@ export function parseToolsFromLog(log: LogEntry): ParsedTool[] { }); // Parse each tool definition - return requestTools.map((tool: ToolDefinition, index: number) => { - const func = tool.function || { name: `Tool ${index + 1}` }; - const name = func.name || `Tool ${index + 1}`; - + // Handle both OpenAI format (tool.function.name) and Anthropic format (tool.name + tool.input_schema) + return requestTools.map((tool: any, index: number) => { + const name = + tool.function?.name || tool.name || `Tool ${index + 1}`; + const description = + tool.function?.description || tool.description || ""; + const parameters = + tool.function?.parameters || tool.input_schema || {}; + return { index: index + 1, - name: name, - description: func.description || "", - parameters: func.parameters || {}, + name, + description, + parameters, called: calledToolNames.has(name), callData: toolCallMap.get(name), }; From 60bcb26dc8a8c064183ca516aff54172bb0b0522 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 24 Feb 2026 18:28:16 -0800 Subject: [PATCH 10/40] feat(agents): assign virtual keys to agents (#22045) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(agents): assign virtual keys to agents - Add agent_id field to LiteLLM_VerificationToken (schema.prisma + _types.py) - Pass agent_id through key generation endpoint so keys can be scoped to an agent - Refactor Add Agent wizard to 3-step flow (Configure → Assign Key → Ready) - Configure: all agent fields, custom/other type with just name+description - Assign Key: create new key or reassign existing key to agent - URL is now optional for easy discovery - Add "Agent" ownership option to Create Key modal on Virtual Keys page with agent selector dropdown - Extract CreatedKeyDisplay into shared component, reused in both flows - Add keyCreateForAgentCall networking helper - Add test for agent_id key generation * fix(agents): code quality fixes from self-review - Fix test_generate_key_helper_fn_agent_id: remove bare except clause, use explicit assert mock_insert.called, use .kwargs for clean arg access - Remove no-op conditional in handleNext (both branches were identical) - Validate selectedExistingKey before calling keyUpdateCall - Validate selectedAgentId before setting on formValues in create_key_button * fix(ui): replace deprecated Tremor Button with Ant Design Button in CreatedKeyDisplay --- litellm/proxy/_types.py | 1 + .../key_management_endpoints.py | 2 + litellm/proxy/schema.prisma | 1 + .../test_key_management_endpoints.py | 54 ++ .../src/components/agents/add_agent_form.tsx | 625 ++++++++++++++---- .../src/components/agents/agent_config.ts | 10 +- .../components/agents/agent_form_fields.tsx | 16 +- .../src/components/networking.tsx | 29 + .../organisms/create_key_button.tsx | 77 ++- .../components/shared/CreatedKeyDisplay.tsx | 53 ++ 10 files changed, 700 insertions(+), 168 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 75b9f91acd9..4053d9d077b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -851,6 +851,7 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): max_budget: Optional[float] = None user_id: Optional[str] = None team_id: Optional[str] = None + agent_id: Optional[str] = None max_parallel_requests: Optional[int] = None metadata: Optional[dict] = {} tpm_limit: Optional[int] = None diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a230c2e9336..c1165ab26d0 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2535,6 +2535,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 user_id: Optional[str] = None, user_alias: Optional[str] = None, team_id: Optional[str] = None, + agent_id: Optional[str] = None, user_email: Optional[str] = None, user_role: Optional[str] = None, max_parallel_requests: Optional[int] = None, @@ -2668,6 +2669,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 "max_budget": key_max_budget, "user_id": user_id, "team_id": team_id, + "agent_id": agent_id, "project_id": project_id, "max_parallel_requests": max_parallel_requests, "metadata": metadata_json, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 23917cf7c7f..8f746b1f9c0 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -314,6 +314,7 @@ model LiteLLM_VerificationToken { router_settings Json? @default("{}") user_id String? team_id String? + agent_id String? project_id String? permissions Json @default("{}") max_parallel_requests Int? diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index ffb4e955426..05df3c2dcbb 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -5623,6 +5623,7 @@ async def test_rotate_master_key_model_data_valid_for_prisma( litellm_params/model_info are JSON strings (create_many expects dicts). """ from unittest.mock import AsyncMock, MagicMock + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.key_management_endpoints import ( _rotate_master_key, @@ -6157,3 +6158,56 @@ async def test_get_member_team_ids(): # Should return team-A and team-B (user is a member of both) # Should NOT return team-C (user is not in members list) assert sorted(result) == ["team-A", "team-B"] + + +@pytest.mark.asyncio +async def test_generate_key_with_agent_id(): + """Test that agent_id is accepted in GenerateKeyRequest and passed to generate_key_helper_fn.""" + from litellm.proxy._types import GenerateKeyRequest + + # Verify GenerateKeyRequest accepts agent_id + request = GenerateKeyRequest( + key_alias="agent-test-key", + agent_id="test-agent-123", + models=[], + ) + assert request.agent_id == "test-agent-123" + data_json = request.model_dump(exclude_unset=True, exclude_none=True) + assert data_json["agent_id"] == "test-agent-123" + + +@pytest.mark.asyncio +async def test_generate_key_helper_fn_agent_id(): + """Test that generate_key_helper_fn passes agent_id into the insert_data call.""" + from unittest.mock import AsyncMock, MagicMock, call, patch + + import litellm.proxy.management_endpoints.key_management_endpoints as km + + mock_prisma_client = AsyncMock() + mock_insert = AsyncMock( + return_value=MagicMock( + token="sk-test", + created_at=None, + updated_at=None, + litellm_budget_table=None, + ) + ) + mock_prisma_client.insert_data = mock_insert + + with patch.object(km, "prisma_client", mock_prisma_client): + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + await generate_key_helper_fn( + request_type="key", + agent_id="test-agent-456", + key_alias="test-agent-key", + models=[], + table_name="key", + ) + + assert mock_insert.called, "insert_data was never called" + # insert_data is called as insert_data(data=key_data, ...) + call_kwargs = mock_insert.call_args.kwargs + key_data = call_kwargs.get("data", {}) + assert key_data.get("agent_id") == "test-agent-456", ( + f"Expected agent_id='test-agent-456' in key_data, got: {key_data.get('agent_id')}" + ) diff --git a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx index f4e0137bd06..360b9b7ee93 100644 --- a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx @@ -1,11 +1,24 @@ import React, { useState, useEffect } from "react"; -import { Modal, Form, message, Select, Input } from "antd"; +import { Modal, Form, message, Select, Input, Steps, Radio, Tag, Divider } from "antd"; import { Button } from "@tremor/react"; -import { createAgentCall, getAgentCreateMetadata, AgentCreateInfo } from "../networking"; +import { CheckCircleFilled, KeyOutlined, RobotOutlined, AppstoreOutlined } from "@ant-design/icons"; +import CreatedKeyDisplay from "../shared/CreatedKeyDisplay"; +import { + createAgentCall, + getAgentCreateMetadata, + keyCreateForAgentCall, + keyListCall, + keyUpdateCall, + AgentCreateInfo, +} from "../networking"; import AgentFormFields from "./agent_form_fields"; import DynamicAgentFormFields, { buildDynamicAgentData } from "./dynamic_agent_form_fields"; import { getDefaultFormValues, buildAgentDataFromForm } from "./agent_config"; +const { Step } = Steps; + +const CUSTOM_AGENT_TYPE = "custom"; + interface AddAgentFormProps { visible: boolean; onClose: () => void; @@ -20,11 +33,25 @@ const AddAgentForm: React.FC = ({ onSuccess, }) => { const [form] = Form.useForm(); + const [currentStep, setCurrentStep] = useState(0); const [isSubmitting, setIsSubmitting] = useState(false); const [agentType, setAgentType] = useState("a2a"); const [agentTypeMetadata, setAgentTypeMetadata] = useState([]); const [loadingMetadata, setLoadingMetadata] = useState(false); + // Step 1: key assignment state + const [keyAssignOption, setKeyAssignOption] = useState<"create_new" | "existing_key" | "skip">("create_new"); + const [newKeyName, setNewKeyName] = useState(""); + const [newKeyModels, setNewKeyModels] = useState([]); + const [existingKeys, setExistingKeys] = useState([]); + const [selectedExistingKey, setSelectedExistingKey] = useState(null); + const [loadingKeys, setLoadingKeys] = useState(false); + + // Step 2: results + const [createdAgentName, setCreatedAgentName] = useState(""); + const [createdKeyValue, setCreatedKeyValue] = useState(null); + const [assignedKeyAlias, setAssignedKeyAlias] = useState(null); + // Fetch agent type metadata on mount useEffect(() => { const fetchMetadata = async () => { @@ -41,11 +68,87 @@ const AddAgentForm: React.FC = ({ fetchMetadata(); }, []); + // Fetch existing keys when assign key step becomes active + useEffect(() => { + if (currentStep === 1 && accessToken && existingKeys.length === 0) { + const fetchKeys = async () => { + setLoadingKeys(true); + try { + const result = await keyListCall(accessToken, null, null, null, null, null, 1, 100); + setExistingKeys(result?.keys || []); + } catch (error) { + console.error("Error fetching keys:", error); + } finally { + setLoadingKeys(false); + } + }; + fetchKeys(); + } + }, [currentStep, accessToken]); + const selectedAgentTypeInfo = agentTypeMetadata.find( (info) => info.agent_type === agentType ); - const handleSubmit = async (values: any) => { + const handleNext = async () => { + try { + if (currentStep === 0) { + await form.validateFields(["agent_name"]); + const agentName = form.getFieldValue("agent_name"); + if (agentName && !newKeyName) { + setNewKeyName(`${agentName}-key`); + } + } + setCurrentStep((s) => s + 1); + } catch { + // validation failed — stay on current step + } + }; + + const handleBack = () => { + setCurrentStep((s) => Math.max(0, s - 1)); + }; + + const buildAgentData = (values: any) => { + if (agentType === CUSTOM_AGENT_TYPE) { + return { + agent_name: values.agent_name, + agent_card_params: { + protocolVersion: "1.0", + name: values.agent_name, + description: values.description || "", + url: "", + version: "1.0.0", + defaultInputModes: ["text"], + defaultOutputModes: ["text"], + capabilities: { streaming: false }, + skills: [], + }, + }; + } else if (agentType === "a2a") { + return buildAgentDataFromForm(values); + } else if (selectedAgentTypeInfo?.use_a2a_form_fields) { + const agentData = buildAgentDataFromForm(values); + if (selectedAgentTypeInfo.litellm_params_template) { + agentData.litellm_params = { + ...agentData.litellm_params, + ...selectedAgentTypeInfo.litellm_params_template, + }; + } + for (const field of selectedAgentTypeInfo.credential_fields) { + const value = values[field.key]; + if (value && field.include_in_litellm_params !== false) { + agentData.litellm_params[field.key] = value; + } + } + return agentData; + } else if (selectedAgentTypeInfo) { + return buildDynamicAgentData(values, selectedAgentTypeInfo); + } + return null; + }; + + const handleCreateAgent = async () => { if (!accessToken) { message.error("No access token available"); return; @@ -53,40 +156,46 @@ const AddAgentForm: React.FC = ({ setIsSubmitting(true); try { - let agentData: any; - - if (agentType === "a2a") { - agentData = buildAgentDataFromForm(values); - } else if (selectedAgentTypeInfo?.use_a2a_form_fields) { - // A2A-compatible agents use the standard A2A form builder - // but need to add litellm_params from the agent type config - agentData = buildAgentDataFromForm(values); - - // Merge litellm_params_template - if (selectedAgentTypeInfo.litellm_params_template) { - agentData.litellm_params = { - ...agentData.litellm_params, - ...selectedAgentTypeInfo.litellm_params_template, - }; - } - - // Add credential fields to litellm_params - for (const field of selectedAgentTypeInfo.credential_fields) { - const value = values[field.key]; - if (value && field.include_in_litellm_params !== false) { - agentData.litellm_params[field.key] = value; - } - } - } else if (selectedAgentTypeInfo) { - agentData = buildDynamicAgentData(values, selectedAgentTypeInfo); + // getFieldsValue(true) returns ALL preserved values including fields from + // unmounted steps; merge with any currently-mounted validated fields. + await form.validateFields(); + const values = { ...form.getFieldsValue(true) }; + const agentData = buildAgentData(values); + if (!agentData) { + message.error("Failed to build agent data"); + setIsSubmitting(false); + return; } - await createAgentCall(accessToken, agentData); - message.success("Agent created successfully"); - form.resetFields(); - setAgentType("a2a"); + const agentResponse = await createAgentCall(accessToken, agentData); + const agentId: string = agentResponse.agent_id; + const agentName: string = agentResponse.agent_name || values.agent_name || agentId; + setCreatedAgentName(agentName); + + if (keyAssignOption === "create_new" && newKeyName) { + const keyResponse = await keyCreateForAgentCall( + accessToken, + agentId, + newKeyName, + newKeyModels, + ); + setCreatedKeyValue(keyResponse.key || null); + } else if (keyAssignOption === "existing_key") { + if (!selectedExistingKey) { + message.error("Please select an existing key to assign"); + setIsSubmitting(false); + return; + } + await keyUpdateCall(accessToken, { + key: selectedExistingKey, + agent_id: agentId, + }); + const keyInfo = existingKeys.find((k) => k.token === selectedExistingKey); + setAssignedKeyAlias(keyInfo?.key_alias || selectedExistingKey.slice(0, 12) + "…"); + } + + setCurrentStep(2); onSuccess(); - onClose(); } catch (error) { console.error("Error creating agent:", error); message.error("Failed to create agent"); @@ -95,9 +204,17 @@ const AddAgentForm: React.FC = ({ } }; - const handleCancel = () => { + const handleClose = () => { form.resetFields(); setAgentType("a2a"); + setCurrentStep(0); + setKeyAssignOption("create_new"); + setNewKeyName(""); + setNewKeyModels([]); + setSelectedExistingKey(null); + setCreatedAgentName(""); + setCreatedKeyValue(null); + setAssignedKeyAlias(null); onClose(); }; @@ -106,25 +223,308 @@ const AddAgentForm: React.FC = ({ form.resetFields(); }; - // Get the logo for the selected agent type for the header - const selectedLogo = selectedAgentTypeInfo?.logo_url || agentTypeMetadata.find(a => a.agent_type === "a2a")?.logo_url; + const isCustomAgent = agentType === CUSTOM_AGENT_TYPE; + const selectedLogo = isCustomAgent + ? null + : selectedAgentTypeInfo?.logo_url || + agentTypeMetadata.find((a) => a.agent_type === "a2a")?.logo_url; + + const renderConfigureStep = () => ( + <> + Agent Type} + required + tooltip="Select the type of agent you want to create" + > + + + +
+ {agentType === CUSTOM_AGENT_TYPE ? ( +
+ + + + + + +
+ ) : agentType === "a2a" ? ( + + ) : selectedAgentTypeInfo?.use_a2a_form_fields ? ( + <> + + {selectedAgentTypeInfo.credential_fields.length > 0 && ( +
+

+ {selectedAgentTypeInfo.agent_type_display_name} Settings +

+ {selectedAgentTypeInfo.credential_fields.map((field) => ( + + {field.field_type === "password" ? ( + + ) : ( + + )} + + ))} +
+ )} + + ) : selectedAgentTypeInfo ? ( + + ) : null} +
+ + ); + + const renderAssignKeyStep = () => { + const agentName = form.getFieldValue("agent_name") || "your-agent"; + return ( +
+ {/* Agent name chip */} +
+ } color="purple" className="px-3 py-1 text-sm"> + {agentName} + +
+ +
+ {/* Option: Create new key */} +
setKeyAssignOption("create_new")} + > +
+
+ setKeyAssignOption("create_new")} + /> +
+
+ + Create a new key for this agent +
+

+ A dedicated key scoped to this agent. +

+ {keyAssignOption === "create_new" && ( +
e.stopPropagation()}> +
+ + setNewKeyName(e.target.value)} + placeholder="e.g. my-agent-key" + /> +
+
+ + setSelectedExistingKey(value)} + filterOption={(input, option) => + (option?.label as string ?? "").toLowerCase().includes(input.toLowerCase()) + } + options={existingKeys.map((k) => ({ + label: k.key_alias || k.token?.slice(0, 12) + "…", + value: k.token, + }))} + /> +
+ )} +
+
+
+
+ +
+ +
+
+ ); + }; + + const renderReadyStep = () => ( +
+ +

Agent Created!

+
+ } color="purple" className="px-3 py-1 text-sm"> + {createdAgentName} + +
+ {createdKeyValue && ( +
+ +
+ )} + {assignedKeyAlias && ( +

+ Key {assignedKeyAlias} has been assigned to this agent. +

+ )} + {!createdKeyValue && !assignedKeyAlias && keyAssignOption === "skip" && ( +

+ No key assigned. You can create one from the Virtual Keys page. +

+ )} +
+ ); return ( - {selectedLogo && ( - Agent + {selectedLogo && currentStep < 1 && ( + Agent )}

Add New Agent

} open={visible} - onCancel={handleCancel} + onCancel={handleClose} footer={null} width={900} className="top-8" @@ -134,103 +534,60 @@ const AddAgentForm: React.FC = ({ }} >
+ {/* Step indicator */} + + + + + +
- {/* Agent Type Selection */} - Agent Type} - required - tooltip="Select the type of agent you want to create" - > - - - - {/* Conditional Form Fields */} -
- {agentType === "a2a" ? ( - - ) : selectedAgentTypeInfo?.use_a2a_form_fields ? ( - // A2A-compatible agents (like Pydantic AI) use full A2A form fields - // plus any additional credential fields - <> - - {selectedAgentTypeInfo.credential_fields.length > 0 && ( -
-

- {selectedAgentTypeInfo.agent_type_display_name} Settings -

- {selectedAgentTypeInfo.credential_fields.map((field) => ( - - {field.field_type === "password" ? ( - - ) : ( - - )} - - ))} -
- )} - - ) : selectedAgentTypeInfo ? ( - - ) : null} -
- - {/* Footer Buttons */} -
- - -
+ {currentStep === 0 && renderConfigureStep()} + {currentStep === 1 && renderAssignKeyStep()} + {currentStep === 2 && renderReadyStep()}
+ + {/* Footer navigation */} +
+
+ {currentStep > 0 && currentStep < 2 && ( + + )} +
+
+ {currentStep < 2 && ( + + )} + {currentStep === 0 && ( + + )} + {currentStep === 1 && ( + + )} + {currentStep === 2 && ( + + )} +
+
); diff --git a/ui/litellm-dashboard/src/components/agents/agent_config.ts b/ui/litellm-dashboard/src/components/agents/agent_config.ts index 9dd41eed8ac..f85c4daac66 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_config.ts +++ b/ui/litellm-dashboard/src/components/agents/agent_config.ts @@ -54,9 +54,9 @@ export const AGENT_FORM_CONFIG: { name: "url", label: "URL", type: "url", - required: true, + required: false, placeholder: "http://localhost:9999/", - tooltip: "Base URL where the agent is hosted", + tooltip: "Base URL where the agent is hosted (optional)", }, { name: "version", @@ -237,9 +237,9 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => { agent_name: values.agent_name, agent_card_params: { protocolVersion: values.protocolVersion || "1.0", - name: values.name, - description: values.description, - url: values.url, + name: values.name || values.agent_name, + description: values.description || "", + url: values.url || "", version: values.version || "1.0.0", defaultInputModes: existingAgent?.agent_card_params?.defaultInputModes || ["text"], defaultOutputModes: existingAgent?.agent_card_params?.defaultOutputModes || ["text"], diff --git a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx index 4dc4ad6829b..d5429d2a3b5 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx @@ -10,13 +10,15 @@ const { Panel } = Collapse; interface AgentFormFieldsProps { showAgentName?: boolean; + visiblePanels?: string[]; } /** * Reusable form fields component for agent forms * Uses shared configuration from agent_config.ts */ -const AgentFormFields: React.FC = ({ showAgentName = true }) => { +const AgentFormFields: React.FC = ({ showAgentName = true, visiblePanels }) => { + const shouldShow = (key: string) => !visiblePanels || visiblePanels.includes(key); return ( <> {showAgentName && ( @@ -32,6 +34,7 @@ const AgentFormFields: React.FC = ({ showAgentName = true {/* Basic Information */} + {shouldShow(AGENT_FORM_CONFIG.basic.key) && ( {AGENT_FORM_CONFIG.basic.fields.map((field) => ( = ({ showAgentName = true ))} + )} {/* Skills */} + {shouldShow(AGENT_FORM_CONFIG.skills.key) && ( {(fields, { add, remove }) => ( @@ -127,8 +132,10 @@ const AgentFormFields: React.FC = ({ showAgentName = true )} + )} {/* Capabilities */} + {shouldShow(AGENT_FORM_CONFIG.capabilities.key) && ( {AGENT_FORM_CONFIG.capabilities.fields.map((field) => ( = ({ showAgentName = true ))} + )} {/* Optional Settings */} + {shouldShow(AGENT_FORM_CONFIG.optional.key) && ( {AGENT_FORM_CONFIG.optional.fields.map((field) => ( = ({ showAgentName = true ))} + )} {/* Cost Configuration */} + {shouldShow(AGENT_FORM_CONFIG.cost.key) && ( + )} {/* LiteLLM Parameters */} + {shouldShow(AGENT_FORM_CONFIG.litellm.key) && ( {AGENT_FORM_CONFIG.litellm.fields.map((field) => ( = ({ showAgentName = true ))} + )} ); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index a3917c2f09e..3aebcb27333 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -930,6 +930,35 @@ export const keyCreateCall = async ( } }; +export const keyCreateForAgentCall = async ( + accessToken: string, + agentId: string, + keyAlias: string, + models: string[], +) => { + const url = proxyBaseUrl ? `${proxyBaseUrl}/key/generate` : `/key/generate`; + const response = await fetch(url, { + method: "POST", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + agent_id: agentId, + key_alias: keyAlias, + models: models.length > 0 ? models : [], + }), + }); + + if (!response.ok) { + const errorData = await response.text(); + handleError(errorData); + throw new Error("Failed to create key for agent"); + } + + return response.json(); +}; + export const userCreateCall = async ( accessToken: string, userID: string | null, diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 9a99870c263..714507da1aa 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -5,7 +5,7 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; import { InfoCircleOutlined } from "@ant-design/icons"; import { useQueryClient } from "@tanstack/react-query"; import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, TextInput, Title } from "@tremor/react"; -import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tooltip } from "antd"; +import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd"; import debounce from "lodash/debounce"; import React, { useCallback, useEffect, useState } from "react"; import { CopyToClipboard } from "react-copy-to-clipboard"; @@ -29,6 +29,7 @@ import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions"; import NotificationsManager from "../molecules/notifications_manager"; import { + getAgentsList, getGuardrailsList, getPoliciesList, getPossibleUserRoles, @@ -42,6 +43,7 @@ import { import NumericalInput from "../shared/numerical_input"; import VectorStoreSelector from "../vector_store_management/VectorStoreSelector"; import { simplifyKeyGenerateError } from "./utils"; +import CreatedKeyDisplay from "../shared/CreatedKeyDisplay"; const { Option } = Select; @@ -169,6 +171,8 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { const [rotationInterval, setRotationInterval] = useState("30d"); const [routerSettings, setRouterSettings] = useState(null); const [routerSettingsKey, setRouterSettingsKey] = useState(0); + const [agentsList, setAgentsList] = useState<{ agent_id: string; agent_name: string }[]>([]); + const [selectedAgentId, setSelectedAgentId] = useState(null); const handleOk = () => { setIsModalVisible(false); form.resetFields(); @@ -180,6 +184,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { setRotationInterval("30d"); setRouterSettings(null); setRouterSettingsKey((prev) => prev + 1); + setSelectedAgentId(null); }; const handleCancel = () => { @@ -195,6 +200,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { setRotationInterval("30d"); setRouterSettings(null); setRouterSettingsKey((prev) => prev + 1); + setSelectedAgentId(null); }; useEffect(() => { @@ -203,6 +209,14 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { } }, [accessToken, userID, userRole]); + useEffect(() => { + if (accessToken) { + getAgentsList(accessToken) + .then((res) => setAgentsList(res?.agents || [])) + .catch(() => setAgentsList([])); + } + }, [accessToken]); + useEffect(() => { const fetchGuardrails = async () => { try { @@ -283,6 +297,12 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { if (keyOwner === "you") { formValues.user_id = userID; + } else if (keyOwner === "agent") { + if (!selectedAgentId) { + message.error("Please select an agent"); + return; + } + formValues.agent_id = selectedAgentId; } // Handle metadata for all key types @@ -539,6 +559,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { You Service Account {userRole === "Admin" && Another User} + Agent New @@ -583,6 +604,32 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => {
)} + {keyOwner === "agent" && ( +
+
+ + Select Agent * + +
+ { + if (filterMode === "single") { + onFiltersChange?.(value ? [value] : []); + } else { + onFiltersChange?.(value); + } + }} options={filterOptions} allowClear /> diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts index ccf0b1b45cf..0f92b61a46e 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts @@ -3,7 +3,7 @@ import type { Team } from "@/components/key_team_helpers/key_list"; export type ExportFormat = "csv" | "json"; export type ExportScope = "daily" | "daily_with_keys" | "daily_with_models"; -export type EntityType = "tag" | "team" | "organization" | "customer" | "agent"; +export type EntityType = "tag" | "team" | "organization" | "customer" | "agent" | "user"; export interface EntitySpendData { results: any[]; diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx index 60674a7c332..c29ade5d653 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx @@ -20,6 +20,7 @@ vi.mock("../../../networking", () => ({ organizationDailyActivityCall: vi.fn(), customerDailyActivityCall: vi.fn(), agentDailyActivityCall: vi.fn(), + userDailyActivityCall: vi.fn(), })); // Mock the child components to simplify testing @@ -58,6 +59,7 @@ describe("EntityUsage", () => { const mockOrganizationDailyActivityCall = vi.mocked(networking.organizationDailyActivityCall); const mockCustomerDailyActivityCall = vi.mocked(networking.customerDailyActivityCall); const mockAgentDailyActivityCall = vi.mocked(networking.agentDailyActivityCall); + const mockUserDailyActivityCall = vi.mocked(networking.userDailyActivityCall); const mockSpendData = { results: [ @@ -146,11 +148,13 @@ describe("EntityUsage", () => { mockOrganizationDailyActivityCall.mockClear(); mockCustomerDailyActivityCall.mockClear(); mockAgentDailyActivityCall.mockClear(); + mockUserDailyActivityCall.mockClear(); mockTagDailyActivityCall.mockResolvedValue(mockSpendData); mockTeamDailyActivityCall.mockResolvedValue(mockSpendData); mockOrganizationDailyActivityCall.mockResolvedValue(mockSpendData); mockCustomerDailyActivityCall.mockResolvedValue(mockSpendData); mockAgentDailyActivityCall.mockResolvedValue(mockSpendData); + mockUserDailyActivityCall.mockResolvedValue(mockSpendData); }); it("should render with tag entity type and display spend metrics", async () => { @@ -232,6 +236,21 @@ describe("EntityUsage", () => { }); }); + it("should render with user entity type and call user API", async () => { + render(); + + await waitFor(() => { + expect(mockUserDailyActivityCall).toHaveBeenCalled(); + }); + + expect(screen.getByText("User Spend Overview")).toBeInTheDocument(); + + await waitFor(() => { + const spendElements = screen.getAllByText("$100.50"); + expect(spendElements.length).toBeGreaterThan(0); + }); + }); + it("should switch between tabs", async () => { render(); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx index 32882341921..a106910cff7 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx @@ -32,6 +32,7 @@ import { organizationDailyActivityCall, tagDailyActivityCall, teamDailyActivityCall, + userDailyActivityCall, } from "../../../networking"; import { getProviderLogoAndName } from "../../../provider_info_helpers"; import { BreakdownMetrics, DailyData, EntityMetricWithMetadata, KeyMetricWithMetadata, TagUsage } from "../../types"; @@ -156,6 +157,15 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti selectedTags.length > 0 ? selectedTags : null, ); setSpendData(data); + } else if (entityType === "user") { + const data = await userDailyActivityCall( + accessToken, + startTime, + endTime, + 1, + selectedTags.length > 0 ? selectedTags[0] : null, + ); + setSpendData(data); } else { throw new Error("Invalid entity type"); } @@ -391,6 +401,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti selectedFilters={selectedTags} onFiltersChange={setSelectedTags} filterOptions={getAllTags() || undefined} + filterMode={entityType === "user" ? "single" : "multiple"} teams={teams || []} /> diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx index d6332e2da80..8b0d5ffac05 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx @@ -6,7 +6,7 @@ * Works at 1m+ spend logs, by querying an aggregate table instead. */ -import { InfoCircleOutlined, LoadingOutlined, UserOutlined } from "@ant-design/icons"; +import { InfoCircleOutlined, LoadingOutlined } from "@ant-design/icons"; import { BarChart, Card, @@ -498,6 +498,36 @@ const UsagePage: React.FC = ({ teams, organizations }) => { {/* Your Usage Panel */} {usageView === "global" && ( <> + {isAdmin && ( +
+ Filter by user + setSelectedUserId(value ?? null)} - filterOption={false} - onSearch={handleUserSearchChange} - searchValue={userSearchInput} - onPopupScroll={handleUserPopupScroll} - loading={isLoadingUsers} - notFoundContent={isLoadingUsers ? : "No users found"} - options={userOptions} - popupRender={(menu) => ( - <> - {menu} - {isFetchingNextUsersPage && ( -
- -
- )} - - )} - /> - {selectedUserId && ( - - Filtering by user - - )} -
- )}
= ({ teams, organizations }) => { dateValue={dateValue} /> )} + {/* User Usage Panel */} + {usageView === "user" && ( + 0 ? userOptions : null} + premiumUser={premiumUser} + dateValue={dateValue} + /> + )} {/* User Agent Activity Panel */} {usageView === "user-agent-activity" && ( diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsageViewSelect/UsageViewSelect.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsageViewSelect/UsageViewSelect.test.tsx index a63fd96fc16..7bb80b424b5 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsageViewSelect/UsageViewSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsageViewSelect/UsageViewSelect.test.tsx @@ -79,6 +79,7 @@ vi.mock("@ant-design/icons", async () => { ShoppingCartOutlined: Icon, TagsOutlined: Icon, RobotOutlined: Icon, + UserOutlined: Icon, LineChartOutlined: Icon, BarChartOutlined: Icon, }; diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsageViewSelect/UsageViewSelect.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsageViewSelect/UsageViewSelect.tsx index 3d456de5a65..53236756d5c 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsageViewSelect/UsageViewSelect.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsageViewSelect/UsageViewSelect.tsx @@ -7,10 +7,11 @@ import { ShoppingCartOutlined, TagsOutlined, TeamOutlined, + UserOutlined, } from "@ant-design/icons"; import { Badge, Select } from "antd"; import React from "react"; -export type UsageOption = "global" | "organization" | "team" | "customer" | "tag" | "agent" | "user-agent-activity"; +export type UsageOption = "global" | "organization" | "team" | "customer" | "tag" | "agent" | "user" | "user-agent-activity"; export interface UsageViewSelectProps { value: UsageOption; onChange: (value: UsageOption) => void; @@ -79,6 +80,13 @@ const OPTIONS: OptionConfig[] = [ icon: , adminOnly: true, }, + { + value: "user", + label: "User Usage", + description: "View usage by individual users", + icon: , + adminOnly: true, + }, { value: "user-agent-activity", label: "User Agent Activity", diff --git a/ui/litellm-dashboard/tsconfig.json b/ui/litellm-dashboard/tsconfig.json index 5b0352feb98..d24bdd340f7 100644 --- a/ui/litellm-dashboard/tsconfig.json +++ b/ui/litellm-dashboard/tsconfig.json @@ -14,7 +14,7 @@ "moduleResolution": "bundler", "resolveJsonModule": true, "isolatedModules": true, - "jsx": "preserve", + "jsx": "react-jsx", "incremental": true, "plugins": [ { From cd41b490614d64e9e429d7eb6f1d18747df374d2 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 25 Feb 2026 10:51:42 -0800 Subject: [PATCH 34/40] fixing build --- .../src/components/ToolPolicies.tsx | 166 +++++++++++------- .../organisms/create_key_button.tsx | 25 +-- 2 files changed, 119 insertions(+), 72 deletions(-) diff --git a/ui/litellm-dashboard/src/components/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx index f69d9cf0c47..860093ceadb 100644 --- a/ui/litellm-dashboard/src/components/ToolPolicies.tsx +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -2,14 +2,7 @@ import React, { useCallback, useDeferredValue, useEffect, useState } from "react"; import { Select, Switch, Tooltip } from "antd"; -import { - Table, - TableHead, - TableHeaderCell, - TableBody, - TableRow, - TableCell, -} from "@tremor/react"; +import { Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell } from "@tremor/react"; import { TimeCell } from "./view_logs/time_cell"; import { TableHeaderSortDropdown } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; import type { SortState } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; @@ -17,14 +10,13 @@ import FilterComponent, { FilterOption } from "./molecules/filter"; import { fetchToolsList, updateToolPolicy, ToolRow } from "./networking"; const POLICY_OPTIONS = [ - { value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" }, - { value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" }, + { value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" }, + { value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" }, ] as const; type PolicyValue = "trusted" | "blocked"; -const policyStyle = (p: string) => - POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1]; +const policyStyle = (p: string) => POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1]; type SortField = "tool_name" | "call_policy" | "team_id" | "key_alias" | "created_at" | "call_count"; @@ -56,18 +48,6 @@ const PolicySelect: React.FC<{ minWidth: 110, fontWeight: 500, }} - styles={{ - selector: { - backgroundColor: style.bg, - borderColor: style.border, - color: style.color, - borderRadius: 999, - fontSize: 11, - fontWeight: 600, - paddingLeft: 8, - paddingRight: 4, - }, - }} popupMatchSelectWidth={false} options={POLICY_OPTIONS.map((o) => ({ value: o.value, @@ -133,7 +113,9 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { } }, [accessToken]); - useEffect(() => { load(); }, [load]); + useEffect(() => { + load(); + }, [load]); useEffect(() => { if (!isLiveTail) return; @@ -146,9 +128,7 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { setSaving(toolName); try { await updateToolPolicy(accessToken, toolName, newPolicy); - setTools((prev) => - prev.map((t) => (t.tool_name === toolName ? { ...t, call_policy: newPolicy } : t)) - ); + setTools((prev) => prev.map((t) => (t.tool_name === toolName ? { ...t, call_policy: newPolicy } : t))); } catch (e: any) { alert(`Failed to update policy: ${e.message}`); } finally { @@ -178,12 +158,14 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { }; // Build unique team/key options from loaded data - const teamOptions = Array.from(new Set(tools.map((t) => t.team_id).filter(Boolean))).map( - (v) => ({ label: v as string, value: v as string }) - ); - const keyAliasOptions = Array.from(new Set(tools.map((t) => t.key_alias).filter(Boolean))).map( - (v) => ({ label: v as string, value: v as string }) - ); + const teamOptions = Array.from(new Set(tools.map((t) => t.team_id).filter(Boolean))).map((v) => ({ + label: v as string, + value: v as string, + })); + const keyAliasOptions = Array.from(new Set(tools.map((t) => t.key_alias).filter(Boolean))).map((v) => ({ + label: v as string, + value: v as string, + })); const filterOptions: FilterOption[] = [ { @@ -245,7 +227,6 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {

Tool Policies

- {/* Toolbar */}
@@ -256,16 +237,29 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { placeholder="Search by Tool Name" className="w-full px-3 py-2 pl-8 border rounded-md text-sm focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500" value={searchTerm} - onChange={(e) => { setSearchTerm(e.target.value); setCurrentPage(1); }} + onChange={(e) => { + setSearchTerm(e.target.value); + setCurrentPage(1); + }} /> - - + +
Live Tail - +
@@ -282,14 +286,27 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
- Showing {filtered.length === 0 ? 0 : (currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, filtered.length)} of {filtered.length} results + Showing {filtered.length === 0 ? 0 : (currentPage - 1) * pageSize + 1} -{" "} + {Math.min(currentPage * pageSize, filtered.length)} of {filtered.length} results + + + Page {currentPage} of {totalPages} - Page {currentPage} of {totalPages}
- - + +
@@ -309,7 +326,9 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { {isLiveTail && (
Auto-refreshing every 15 seconds - +
)} @@ -321,20 +340,34 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { - - - - - + + + + + + + + + + + + + + + Key Hash - + + + Origin {loading ? ( - Loading tools… + + Loading tools… + ) : paginated.length === 0 ? ( @@ -397,12 +430,25 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { {/* Bottom pagination (only when > 1 page) */} {totalPages > 1 && (
- Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of {sorted.length} + + Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of{" "} + {sorted.length} +
- - + +
)} diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 714507da1aa..fc80c394ae6 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -8,10 +8,10 @@ import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, Tex import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd"; import debounce from "lodash/debounce"; import React, { useCallback, useEffect, useState } from "react"; -import { CopyToClipboard } from "react-copy-to-clipboard"; import { rolesWithWriteAccess } from "../../utils/roles"; import AgentSelector from "../agent_management/AgentSelector"; import { mapDisplayToInternalNames } from "../callback_info_helpers"; +import AccessGroupSelector from "../common_components/AccessGroupSelector"; import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"; import SchemaFormFields from "../common_components/check_openapi_schema"; import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings"; @@ -20,7 +20,6 @@ import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSel import PremiumLoggingSettings from "../common_components/PremiumLoggingSettings"; import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"; import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "../common_components/RouterSettingsAccordion"; -import AccessGroupSelector from "../common_components/AccessGroupSelector"; import TeamDropdown from "../common_components/team_dropdown"; import { CreateUserButton } from "../CreateUserButton"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; @@ -40,10 +39,10 @@ import { proxyBaseUrl, userFilterUICall, } from "../networking"; +import CreatedKeyDisplay from "../shared/CreatedKeyDisplay"; import NumericalInput from "../shared/numerical_input"; import VectorStoreSelector from "../vector_store_management/VectorStoreSelector"; import { simplifyKeyGenerateError } from "./utils"; -import CreatedKeyDisplay from "../shared/CreatedKeyDisplay"; const { Option } = Select; @@ -299,7 +298,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { formValues.user_id = userID; } else if (keyOwner === "agent") { if (!selectedAgentId) { - message.error("Please select an agent"); + NotificationsManager.error("Please select an agent"); return; } formValues.agent_id = selectedAgentId; @@ -559,7 +558,9 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { You Service Account {userRole === "Admin" && Another User} - Agent New + + Agent New + @@ -1005,9 +1006,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { style={{ width: "100%" }} disabled={!premiumUser} placeholder={ - !premiumUser - ? "Premium feature - Upgrade to set policies by key" - : "Select or enter policies" + !premiumUser ? "Premium feature - Upgrade to set policies by key" : "Select or enter policies" } options={policiesList.map((name) => ({ value: name, label: name }))} /> @@ -1059,9 +1058,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { className="mt-4" help="Select access groups to assign to this key" > - + = ({ team, teams, data, addKey }) => { accessToken={accessToken || ""} value={routerSettings || undefined} onChange={setRouterSettings} - modelData={userModels.length > 0 ? { data: userModels.map((model) => ({ model_name: model })) } : undefined} + modelData={ + userModels.length > 0 + ? { data: userModels.map((model) => ({ model_name: model })) } + : undefined + } /> From 7daeaf81063cf45859bfc7028fc610d0b1f55a8c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 25 Feb 2026 10:52:24 -0800 Subject: [PATCH 35/40] [Docs] Add Credential Usage Tracking documentation Add new document explaining automatic credential usage tracking and tagging. When models use reusable credentials, LiteLLM automatically injects a Credential: tag on requests, enabling credential-level spend tracking on the Usage page with no additional configuration. Co-Authored-By: Claude Haiku 4.5 --- .../docs/proxy/credential_usage_tracking.md | 19 +++++++++++++++++++ docs/my-website/docs/proxy/ui_credentials.md | 4 ++++ 2 files changed, 23 insertions(+) create mode 100644 docs/my-website/docs/proxy/credential_usage_tracking.md diff --git a/docs/my-website/docs/proxy/credential_usage_tracking.md b/docs/my-website/docs/proxy/credential_usage_tracking.md new file mode 100644 index 00000000000..25658144c49 --- /dev/null +++ b/docs/my-website/docs/proxy/credential_usage_tracking.md @@ -0,0 +1,19 @@ +# Credential Usage Tracking + +When a model is attached to a [reusable credential](./ui_credentials.md), LiteLLM automatically injects the credential name as a tag on every request that uses that model. This means credential-level spend and usage are tracked with zero extra configuration. + +## How It Works + +When you attach a model to a reusable credential via `litellm_credential_name`, each request routed through that model is tagged `Credential: ` (for example, `Credential: xAI`). This tag flows into `DailyTagSpend` and appears in the **Tag** view on the Usage page, where you can filter spend and usage by credential. + +If a model has no credential attached, behavior is unchanged—no credential tag is added. + +## Viewing Credential Usage + +In the Admin UI, go to **Usage → Tag** and look for tags with the `Credential: ` prefix. These represent aggregated spend and token usage across all requests that used that credential. + +## Related Documentation + +- [Adding LLM Credentials](./ui_credentials.md) - How to create and attach reusable credentials to models +- [Tag Budgets](./tag_budgets.md) - Setting spend limits on tags +- [Tag Routing](./tag_routing.md) - Routing requests based on tags diff --git a/docs/my-website/docs/proxy/ui_credentials.md b/docs/my-website/docs/proxy/ui_credentials.md index 40db5368596..f10f2631f83 100644 --- a/docs/my-website/docs/proxy/ui_credentials.md +++ b/docs/my-website/docs/proxy/ui_credentials.md @@ -46,6 +46,10 @@ Go to Add Model -> Existing Credentials -> Select your credential in the dropdow +## Usage Tracking + +Models attached to a reusable credential are automatically tracked in the Usage page. Each request is tagged `Credential: ` and appears in the **Tag** view, so you can filter spend and usage by credential without any extra configuration. See [Credential Usage Tracking](./credential_usage_tracking.md) for details. + ## Frequently Asked Questions From 30e8151288e32042406841220164360225ba79e0 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 25 Feb 2026 11:30:09 -0800 Subject: [PATCH 36/40] fixing build and tests --- .../components/UsagePage/components/UsagePageView.test.tsx | 4 ++-- .../src/components/organisms/create_key_button.test.tsx | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx index 9a603252286..410d7510171 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx @@ -713,7 +713,7 @@ describe("UsagePage", () => { // Admin should see the user selector select element with the placeholder attribute const userSelects = screen.getAllByRole("combobox"); const userSelect = userSelects.find( - (el) => el.getAttribute("placeholder") === "All Users (Global View)", + (el) => el.getAttribute("placeholder") === "Select user to filter...", ); expect(userSelect).toBeDefined(); }); @@ -828,7 +828,7 @@ describe("UsagePage", () => { // Non-admin should not see the user selector const userSelects = screen.getAllByRole("combobox"); const userSelect = userSelects.find( - (el) => el.getAttribute("placeholder") === "All Users (Global View)", + (el) => el.getAttribute("placeholder") === "Select user to filter...", ); expect(userSelect).toBeUndefined(); }); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx index 46dca6039a7..12118a6aaa0 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx @@ -27,6 +27,7 @@ vi.mock("../networking", () => ({ soft_budget: null, }), fetchMCPAccessGroups: vi.fn().mockResolvedValue([]), + getAgentsList: vi.fn().mockResolvedValue([]), })); vi.mock("../molecules/notifications_manager", () => ({ From 12c4876891dc0ed81560700f764516201bac811b Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Wed, 25 Feb 2026 11:44:30 -0800 Subject: [PATCH 37/40] Agents - assign tools (#22064) * feat(proxy): add max_iterations limiter for agent session loops (#22058) Adds a new proxy hook that enforces a per-session cap on the number of LLM calls an agentic loop can make. Callers send a session_id with each request, and the hook counts calls per session, returning 429 when the configured max_iterations limit is exceeded. - Uses Redis Lua script for atomic increment (multi-instance safe) - Falls back to in-memory cache when Redis unavailable - Follows parallel_request_limiter_v3 pattern - Configurable via key metadata: {"max_iterations": 25} - Session counters auto-expire via TTL (default 1hr) Co-authored-by: Claude Opus 4.6 * feat: add new code execution dataset * feat(agent_endpoints/): allow giving agents keys * fix: ui fixes * feat: allow assigning mcp servers to agents * fix: eliminate duplicate DB queries in MCP agent auth and N+1 in agent listing (#22110) - Extract _get_agent_object_permission helper so _get_allowed_mcp_servers_for_agent and _get_agent_tool_permissions_for_server share a single DB fetch instead of each independently querying the same agent row (was 1+N queries per MCP request) - Use include={"object_permission": True} on find_many in get_all_agents_from_db to eagerly load permissions in one query instead of N+1 - Use include={"object_permission": True} on create/update/find_unique in all agent CRUD operations, removing attach_object_permission_to_dict follow-up calls Co-authored-by: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- .../migration.sql | 40 ++++ .../litellm_proxy_extras/schema.prisma | 25 ++- .../mcp_server/auth/user_api_key_auth_mcp.py | 171 +++++++++++++- litellm/proxy/_types.py | 76 +++---- .../proxy/agent_endpoints/agent_registry.py | 121 ++++++++-- litellm/proxy/agent_endpoints/endpoints.py | 18 +- litellm/proxy/hooks/max_iterations_limiter.py | 208 ++++++++++++++++++ litellm/types/agents.py | 10 + schema.prisma | 25 ++- scripts/test_agent_mcp_endpoints.sh | 186 ++++++++++++++++ .../auth/test_user_api_key_auth_mcp.py | 143 ++++++++++++ .../hooks/test_max_iterations_limiter.py | 106 +++++++++ ui/litellm-dashboard/package-lock.json | 15 ++ .../components/ModelRetrySettingsTab.test.tsx | 2 +- .../src/components/agents.tsx | 62 +++++- .../src/components/agents/add_agent_form.tsx | 145 ++++++++++-- .../src/components/agents/agent_card.tsx | 103 +++++++++ .../src/components/agents/agent_card_grid.tsx | 63 ++++++ .../src/components/agents/agent_info.tsx | 38 ++++ .../src/components/agents/types.ts | 14 ++ .../guardrails/guardrail_garden.tsx | 2 +- .../mcp_tools/mcp_tool_configuration.tsx | 2 +- .../src/components/mcp_tools/mcp_tools.tsx | 10 +- .../src/components/networking.tsx | 2 + .../organisms/create_key_button.tsx | 2 +- .../policies/pipeline_flow_builder.tsx | 2 +- .../components/policies/policy_templates.tsx | 2 +- .../templates/KeyInfoHeader.test.tsx | 2 +- .../LogDetailsDrawer/LogDetailContent.tsx | 1 - .../src/components/view_logs/index.tsx | 2 +- .../src/data/financialCompliancePrompts.ts | 2 +- .../src/data/insultsCompliancePrompts.ts | 2 +- 32 files changed, 1487 insertions(+), 115 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260224203854_add_agent_object_permissions_table/migration.sql create mode 100644 litellm/proxy/hooks/max_iterations_limiter.py create mode 100755 scripts/test_agent_mcp_endpoints.sh create mode 100644 tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py create mode 100644 ui/litellm-dashboard/src/components/agents/agent_card.tsx create mode 100644 ui/litellm-dashboard/src/components/agents/agent_card_grid.tsx diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224203854_add_agent_object_permissions_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224203854_add_agent_object_permissions_table/migration.sql new file mode 100644 index 00000000000..78e364d5478 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224203854_add_agent_object_permissions_table/migration.sql @@ -0,0 +1,40 @@ +-- AlterTable +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "object_permission_id" TEXT; + +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN "spec_path"; + +-- AlterTable +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "agent_id" TEXT; + +-- CreateTable +CREATE TABLE "LiteLLM_ToolTable" ( + "tool_id" TEXT NOT NULL, + "tool_name" TEXT NOT NULL, + "origin" TEXT, + "call_policy" TEXT NOT NULL DEFAULT 'untrusted', + "call_count" INTEGER NOT NULL DEFAULT 0, + "assignments" JSONB DEFAULT '{}', + "key_hash" TEXT, + "team_id" TEXT, + "key_alias" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT, + + CONSTRAINT "LiteLLM_ToolTable_pkey" PRIMARY KEY ("tool_id") +); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_ToolTable_tool_name_key" ON "LiteLLM_ToolTable"("tool_name"); + +-- CreateIndex +CREATE INDEX "LiteLLM_ToolTable_call_policy_idx" ON "LiteLLM_ToolTable"("call_policy"); + +-- CreateIndex +CREATE INDEX "LiteLLM_ToolTable_team_id_idx" ON "LiteLLM_ToolTable"("team_id"); + +-- AddForeignKey +ALTER TABLE "LiteLLM_AgentsTable" ADD CONSTRAINT "LiteLLM_AgentsTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 4af7484148c..155cea12ca4 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -64,6 +64,8 @@ model LiteLLM_AgentsTable { litellm_params Json? agent_card_params Json agent_access_groups String[] @default([]) + object_permission_id String? + object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@ -264,6 +266,7 @@ model LiteLLM_ObjectPermissionTable { organizations LiteLLM_OrganizationTable[] users LiteLLM_UserTable[] end_users LiteLLM_EndUserTable[] + agents_table LiteLLM_AgentsTable[] } // Holds the MCP server configuration @@ -273,7 +276,6 @@ model LiteLLM_MCPServerTable { alias String? description String? url String? - spec_path String? transport String @default("sse") auth_type String? credentials Json? @default("{}") @@ -315,6 +317,7 @@ model LiteLLM_VerificationToken { router_settings Json? @default("{}") user_id String? team_id String? + agent_id String? project_id String? permissions Json @default("{}") max_parallel_requests Int? @@ -1052,6 +1055,26 @@ model LiteLLM_PolicyAttachmentTable { updated_by String? } +// Global tool registry - auto-discovered from LLM responses; admins set call_policy here +model LiteLLM_ToolTable { + tool_id String @id @default(uuid()) + tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space" + origin String? // MCP server name or "user_defined" + call_policy String @default("untrusted") // "trusted" | "untrusted" | "dual_llm" | "blocked" + call_count Int @default(0) // cumulative number of times this tool was seen + assignments Json? @default("{}") + key_hash String? // hash of the virtual key that first called this tool + team_id String? // team that first called this tool + key_alias String? // human-readable alias of the virtual key + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? + + @@index([call_policy]) + @@index([team_id]) +} + //Unified Access Groups table for storing unified access groups model LiteLLM_AccessGroupTable { access_group_id String @id @default(uuid()) 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 60b29b975f7..860569d24cb 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 @@ -412,6 +412,26 @@ class MCPRequestHandler: ) return [] + ######################################################### + # Check agent permissions if agent_id is set on the key + ######################################################### + if user_api_key_auth and user_api_key_auth.agent_id: + allowed_mcp_servers_for_agent = ( + await MCPRequestHandler._get_allowed_mcp_servers_for_agent( + user_api_key_auth + ) + ) + if len(allowed_mcp_servers_for_agent) > 0: + # Intersect: agent can only use servers allowed by BOTH key/team AND agent config + allowed_mcp_servers = [ + s + for s in allowed_mcp_servers + if s in allowed_mcp_servers_for_agent + ] + verbose_logger.debug( + f"Applied agent intersection filter. Final allowed servers: {allowed_mcp_servers}" + ) + return list(set(allowed_mcp_servers)) except Exception as e: verbose_logger.warning(f"Failed to get allowed MCP servers: {str(e)}") @@ -513,13 +533,33 @@ class MCPRequestHandler: if team_tools: if key_tools: # Both have restrictions → intersection - return list(set(team_tools) & set(key_tools)) + allowed_tools = list(set(team_tools) & set(key_tools)) else: # Only team has restrictions → inherit from team - return team_tools + allowed_tools = team_tools else: # No team restrictions → use key restrictions - return key_tools + allowed_tools = key_tools + + # Intersect with agent's tool permissions if agent_id is set + if user_api_key_auth.agent_id: + # Pre-fetch agent object_permission once to avoid duplicate DB query + agent_obj_perm = await MCPRequestHandler._get_agent_object_permission( + user_api_key_auth + ) + agent_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( + server_id=server_id, + user_api_key_auth=user_api_key_auth, + agent_object_permission=agent_obj_perm, + ) + if agent_tools is not None: + if allowed_tools is not None: + allowed_tools = list( + set(allowed_tools) & set(agent_tools) + ) + else: + allowed_tools = agent_tools + return allowed_tools except Exception as e: verbose_logger.warning(f"Failed to get allowed tools for server: {str(e)}") @@ -715,6 +755,131 @@ class MCPRequestHandler: ) return [] + @staticmethod + async def _get_agent_object_permission( + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + ): + """ + Fetch the agent's object_permission from the DB (single query). + + Returns the object_permission object or None. + """ + from litellm.proxy.proxy_server import prisma_client + + if not user_api_key_auth or not user_api_key_auth.agent_id: + return None + + if prisma_client is None: + verbose_logger.debug("prisma_client is None") + return None + + try: + agent_row = await prisma_client.db.litellm_agentstable.find_unique( + where={"agent_id": user_api_key_auth.agent_id}, + include={"object_permission": True}, + ) + if agent_row is None or agent_row.object_permission is None: + return None + + return agent_row.object_permission + except Exception as e: + verbose_logger.warning( + f"Failed to get agent object permission: {str(e)}" + ) + return None + + @staticmethod + async def _get_allowed_mcp_servers_for_agent( + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + agent_object_permission=None, + ) -> List[str]: + """ + Get allowed MCP servers for an agent (from the agent's object_permission). + + Returns the MCP servers from the agent's object_permission. + If agent has no object_permission, returns [] (no extra restriction). + + Args: + user_api_key_auth: User auth with agent_id + agent_object_permission: Pre-fetched object_permission to avoid duplicate DB query. + If None, will be fetched from DB. + """ + if not user_api_key_auth or not user_api_key_auth.agent_id: + return [] + + try: + obj_perm = agent_object_permission + if obj_perm is None: + obj_perm = await MCPRequestHandler._get_agent_object_permission( + user_api_key_auth + ) + if obj_perm is None: + return [] + + direct_mcp_servers = getattr(obj_perm, "mcp_servers", None) or [] + if isinstance(direct_mcp_servers, str): + direct_mcp_servers = [] + mcp_access_groups = getattr(obj_perm, "mcp_access_groups", None) or [] + if isinstance(mcp_access_groups, str): + mcp_access_groups = [] + + access_group_servers = ( + await MCPRequestHandler._get_mcp_servers_from_access_groups( + mcp_access_groups + ) + ) + all_servers = list(direct_mcp_servers) + access_group_servers + return list(set(all_servers)) + except Exception as e: + verbose_logger.warning( + f"Failed to get allowed MCP servers for agent: {str(e)}" + ) + return [] + + @staticmethod + async def _get_agent_tool_permissions_for_server( + server_id: str, + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + agent_object_permission=None, + ) -> Optional[List[str]]: + """ + Get allowed tool names for a server from the agent's object_permission. + Returns None if agent has no tool restrictions for this server. + + Args: + server_id: Server ID to check permissions for + user_api_key_auth: User auth with agent_id + agent_object_permission: Pre-fetched object_permission to avoid duplicate DB query. + If None, will be fetched from DB. + """ + if not user_api_key_auth or not user_api_key_auth.agent_id: + return None + + try: + obj_perm = agent_object_permission + if obj_perm is None: + obj_perm = await MCPRequestHandler._get_agent_object_permission( + user_api_key_auth + ) + if obj_perm is None: + return None + + mcp_tool_permissions = getattr( + obj_perm, "mcp_tool_permissions", None + ) + if not mcp_tool_permissions: + return None + if isinstance(mcp_tool_permissions, dict): + tools = mcp_tool_permissions.get(server_id) + else: + tools = None + return list(tools) if tools else None + except Exception as e: + verbose_logger.warning( + f"Failed to get agent tool permissions for server: {str(e)}" + ) + return None + @staticmethod def _get_config_server_ids_for_access_groups( config_mcp_servers, access_groups: List[str] diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4053d9d077b..5430d7a3605 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,60 +1,40 @@ import enum import json from datetime import datetime -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union +from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal, + Optional, Union) import httpx -from pydantic import ( - BaseModel, - ConfigDict, - Field, - Json, - field_validator, - model_validator, -) +from pydantic import (BaseModel, ConfigDict, Field, Json, field_validator, + model_validator) from typing_extensions import Required, TypedDict from litellm._uuid import uuid from litellm.types.integrations.slack_alerting import AlertType -from litellm.types.llms.openai import ( - AllMessageValues, - OpenAIFileObject, - ResponsesAPIResponse, -) -from litellm.types.mcp import ( - MCPAuth, - MCPAuthType, - MCPCredentials, - MCPTransport, - MCPTransportType, -) +from litellm.types.llms.openai import (AllMessageValues, OpenAIFileObject, + ResponsesAPIResponse) +from litellm.types.mcp import (MCPAuth, MCPAuthType, MCPCredentials, + MCPTransport, MCPTransportType) from litellm.types.mcp_server.mcp_server_manager import MCPInfo from litellm.types.router import RouterErrors, UpdateRouterConfig from litellm.types.secret_managers.main import KeyManagementSystem -from litellm.types.utils import ( - CallTypes, - CostBreakdown, - EmbeddingResponse, - GenericBudgetConfigType, - ImageResponse, - LiteLLMBatch, - LiteLLMFineTuningJob, - LiteLLMPydanticObjectBase, - ModelResponse, - ProviderField, - StandardCallbackDynamicParams, - StandardLoggingGuardrailInformation, - StandardLoggingMCPToolCall, - StandardLoggingModelInformation, - StandardLoggingPayloadErrorInformation, - StandardLoggingPayloadStatus, - StandardLoggingVectorStoreRequest, - StandardPassThroughResponseObject, - TextCompletionResponse, -) +from litellm.types.utils import (CallTypes, CostBreakdown, EmbeddingResponse, + GenericBudgetConfigType, ImageResponse, + LiteLLMBatch, LiteLLMFineTuningJob, + LiteLLMPydanticObjectBase, ModelResponse, + ProviderField, StandardCallbackDynamicParams, + StandardLoggingGuardrailInformation, + StandardLoggingMCPToolCall, + StandardLoggingModelInformation, + StandardLoggingPayloadErrorInformation, + StandardLoggingPayloadStatus, + StandardLoggingVectorStoreRequest, + StandardPassThroughResponseObject, + TextCompletionResponse) from litellm.types.videos.main import VideoObject -from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type +from .types_utils.utils import (get_instance_fn, + validate_custom_validate_return_type) if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -2202,6 +2182,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): config: Dict = {} user_id: Optional[str] = None team_id: Optional[str] = None + agent_id: Optional[str] = None project_id: Optional[str] = None max_parallel_requests: Optional[int] = None metadata: Dict = {} @@ -2379,7 +2360,8 @@ class UserAPIKeyAuth( This is used to track number of requests/spend for health check calls. """ - from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME + from litellm.constants import \ + LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME return cls( api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, @@ -2411,7 +2393,8 @@ class UserAPIKeyAuth( This is used to track actions performed by automated system jobs. """ - from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + from litellm.constants import \ + LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME return cls( api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, @@ -2802,7 +2785,8 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): @model_validator(mode="after") def mask_api_keys(self): - from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker + from litellm.litellm_core_utils.sensitive_data_masker import \ + SensitiveDataMasker masker = SensitiveDataMasker(sensitive_patterns={"key"}) diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 0d2df3856a1..159c9fb93d9 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -5,6 +5,9 @@ from typing import Any, Dict, List, Optional import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy.management_helpers.object_permission_utils import ( + handle_update_object_permission_common, +) from litellm.proxy.utils import PrismaClient from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -117,20 +120,39 @@ class AgentRegistry: ) agent_card_params: str = safe_dumps(agent_card_params_dict) + # Handle object_permission (MCP tool access for agent) + object_permission_id: Optional[str] = None + if agent.get("object_permission") is not None: + agent_copy = dict(agent) + object_permission_id = await handle_update_object_permission_common( + agent_copy, None, prisma_client + ) + + create_data: Dict[str, Any] = { + "agent_name": agent_name, + "litellm_params": litellm_params, + "agent_card_params": agent_card_params, + "created_by": created_by, + "updated_by": created_by, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + } + if object_permission_id is not None: + create_data["object_permission_id"] = object_permission_id + # Create agent in DB created_agent = await prisma_client.db.litellm_agentstable.create( - data={ - "agent_name": agent_name, - "litellm_params": litellm_params, - "agent_card_params": agent_card_params, - "created_by": created_by, - "updated_by": created_by, - "created_at": datetime.now(timezone.utc), - "updated_at": datetime.now(timezone.utc), - } + data=create_data, + include={"object_permission": True}, ) - return AgentResponse(**created_agent.model_dump()) # type: ignore + created_agent_dict = created_agent.model_dump() + if created_agent.object_permission is not None: + try: + created_agent_dict["object_permission"] = created_agent.object_permission.model_dump() + except Exception: + created_agent_dict["object_permission"] = created_agent.object_permission.dict() + return AgentResponse(**created_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error adding agent to DB: {str(e)}") @@ -181,7 +203,7 @@ class AgentRegistry: raise Exception(f"Agent with ID {agent_id} not found") augment_agent = {**existing_agent, **agent} - update_data = {} + update_data: Dict[str, Any] = {} if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") if augment_agent.get("litellm_params"): @@ -192,6 +214,20 @@ class AgentRegistry: update_data["agent_card_params"] = safe_dumps( augment_agent.get("agent_card_params") ) + if agent.get("object_permission") is not None: + agent_copy = dict(augment_agent) + existing_object_permission_id = existing_agent.get( + "object_permission_id" + ) + object_permission_id = ( + await handle_update_object_permission_common( + agent_copy, + existing_object_permission_id, + prisma_client, + ) + ) + if object_permission_id is not None: + update_data["object_permission_id"] = object_permission_id # Patch agent in DB patched_agent = await prisma_client.db.litellm_agentstable.update( where={"agent_id": agent_id}, @@ -200,8 +236,15 @@ class AgentRegistry: "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), }, + include={"object_permission": True}, ) - return AgentResponse(**patched_agent.model_dump()) # type: ignore + patched_agent_dict = patched_agent.model_dump() + if patched_agent.object_permission is not None: + try: + patched_agent_dict["object_permission"] = patched_agent.object_permission.model_dump() + except Exception: + patched_agent_dict["object_permission"] = patched_agent.object_permission.dict() + return AgentResponse(**patched_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error patching agent in DB: {str(e)}") @@ -238,19 +281,47 @@ class AgentRegistry: ) agent_card_params: str = safe_dumps(agent_card_params_dict) + update_data: Dict[str, Any] = { + "agent_name": agent_name, + "litellm_params": litellm_params, + "agent_card_params": agent_card_params, + "updated_by": updated_by, + "updated_at": datetime.now(timezone.utc), + } + if agent.get("object_permission") is not None: + existing_agent = await prisma_client.db.litellm_agentstable.find_unique( + where={"agent_id": agent_id} + ) + existing_object_permission_id = ( + existing_agent.object_permission_id + if existing_agent is not None + else None + ) + agent_copy = dict(agent) + object_permission_id = ( + await handle_update_object_permission_common( + agent_copy, + existing_object_permission_id, + prisma_client, + ) + ) + if object_permission_id is not None: + update_data["object_permission_id"] = object_permission_id + # Update agent in DB updated_agent = await prisma_client.db.litellm_agentstable.update( where={"agent_id": agent_id}, - data={ - "agent_name": agent_name, - "litellm_params": litellm_params, - "agent_card_params": agent_card_params, - "updated_by": updated_by, - "updated_at": datetime.now(timezone.utc), - }, + data=update_data, + include={"object_permission": True}, ) - return AgentResponse(**updated_agent.model_dump()) # type: ignore + updated_agent_dict = updated_agent.model_dump() + if updated_agent.object_permission is not None: + try: + updated_agent_dict["object_permission"] = updated_agent.object_permission.model_dump() + except Exception: + updated_agent_dict["object_permission"] = updated_agent.object_permission.dict() + return AgentResponse(**updated_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error updating agent in DB: {str(e)}") @@ -264,11 +335,19 @@ class AgentRegistry: try: agents_from_db = await prisma_client.db.litellm_agentstable.find_many( order={"created_at": "desc"}, + include={"object_permission": True}, ) agents: List[Dict[str, Any]] = [] for agent in agents_from_db: - agents.append(dict(agent)) + agent_dict = dict(agent) + # object_permission is eagerly loaded via include above + if agent.object_permission is not None: + try: + agent_dict["object_permission"] = agent.object_permission.model_dump() + except Exception: + agent_dict["object_permission"] = agent.object_permission.dict() + agents.append(agent_dict) return agents except Exception as e: diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 4a8d615f0b3..b411b81b434 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -16,6 +16,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity from litellm.types.agents import ( AgentConfig, AgentMakePublicResponse, @@ -23,8 +24,6 @@ from litellm.types.agents import ( MakeAgentsPublicRequest, PatchAgentRequest, ) - -from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) @@ -233,11 +232,18 @@ async def get_agent_by_id(agent_id: str): try: agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: - agent = await prisma_client.db.litellm_agentstable.find_unique( - where={"agent_id": agent_id} + agent_row = await prisma_client.db.litellm_agentstable.find_unique( + where={"agent_id": agent_id}, + include={"object_permission": True}, ) - if agent is not None: - agent = AgentResponse(**agent.model_dump()) # type: ignore + if agent_row is not None: + agent_dict = agent_row.model_dump() + if agent_row.object_permission is not None: + try: + agent_dict["object_permission"] = agent_row.object_permission.model_dump() + except Exception: + agent_dict["object_permission"] = agent_row.object_permission.dict() + agent = AgentResponse(**agent_dict) # type: ignore if agent is None: raise HTTPException( diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py new file mode 100644 index 00000000000..8d481f6b261 --- /dev/null +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -0,0 +1,208 @@ +""" +Max Iterations Limiter for LiteLLM Proxy. + +Enforces a per-session cap on the number of LLM calls an agentic loop can make. +Callers send a `session_id` with each request (via `x-litellm-session-id` header +or `metadata.session_id`), and this hook counts calls per session. When the count +exceeds `max_iterations` (configured in key/team metadata), returns 429. + +Works across multiple proxy instances via DualCache (in-memory + Redis). +Follows the same pattern as parallel_request_limiter_v3.py. +""" + +import os +from typing import TYPE_CHECKING, Any, Optional, Union + +from fastapi import HTTPException + +from litellm import DualCache +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth + +if TYPE_CHECKING: + from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache + + InternalUsageCache = _InternalUsageCache +else: + InternalUsageCache = Any + + +# Redis Lua script for atomic increment with TTL. +# Returns the new count after increment. +# Only sets EXPIRE on first increment (when count becomes 1). +MAX_ITERATIONS_INCREMENT_SCRIPT = """ +local key = KEYS[1] +local ttl = tonumber(ARGV[1]) + +local current = redis.call('INCR', key) +if current == 1 then + redis.call('EXPIRE', key, ttl) +end + +return current +""" + +# Default TTL for session iteration counters (1 hour) +DEFAULT_MAX_ITERATIONS_TTL = 3600 + + +class _PROXY_MaxIterationsHandler(CustomLogger): + """ + Pre-call hook that enforces max_iterations per session. + + Configuration: + - max_iterations: set in key metadata via /key/generate or /key/update + e.g. metadata={"max_iterations": 25} + - session_id: sent by caller via x-litellm-session-id header or + metadata.session_id in request body + + Cache key pattern: + {session_iterations:}:count + + Multi-instance support: + Uses Redis Lua script for atomic increment (same pattern as + parallel_request_limiter_v3). Falls back to in-memory cache + when Redis is unavailable. + """ + + def __init__(self, internal_usage_cache: InternalUsageCache): + self.internal_usage_cache = internal_usage_cache + self.ttl = int( + os.getenv("LITELLM_MAX_ITERATIONS_TTL", DEFAULT_MAX_ITERATIONS_TTL) + ) + + # Register Lua script with Redis if available (same pattern as v3 limiter) + if self.internal_usage_cache.dual_cache.redis_cache is not None: + self.increment_script = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + MAX_ITERATIONS_INCREMENT_SCRIPT + ) + ) + else: + self.increment_script = None + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: str, + ) -> Optional[Union[Exception, str, dict]]: + """ + Check session iteration count before making the API call. + + Extracts session_id from request metadata and max_iterations from + key metadata. If the session has exceeded max_iterations, raises 429. + """ + # Extract session_id from request data + session_id = self._get_session_id(data) + if session_id is None: + return None + + # Extract max_iterations from key metadata + max_iterations = self._get_max_iterations(user_api_key_dict) + if max_iterations is None: + return None + + verbose_proxy_logger.debug( + "MaxIterationsHandler: session_id=%s, max_iterations=%s", + session_id, + max_iterations, + ) + + # Increment and check + cache_key = self._make_cache_key(session_id) + current_count = await self._increment_and_get(cache_key) + + if current_count > max_iterations: + raise HTTPException( + status_code=429, + detail=( + f"Max iterations exceeded for session {session_id}. " + f"Current count: {current_count}, max_iterations: {max_iterations}." + ), + ) + + verbose_proxy_logger.debug( + "MaxIterationsHandler: session_id=%s, count=%s/%s", + session_id, + current_count, + max_iterations, + ) + + return None + + def _get_session_id(self, data: dict) -> Optional[str]: + """Extract session_id from request metadata.""" + metadata = data.get("metadata") or {} + session_id = metadata.get("session_id") + if session_id is not None: + return str(session_id) + + # Also check litellm_metadata (used for /thread and /assistant endpoints) + litellm_metadata = data.get("litellm_metadata") or {} + session_id = litellm_metadata.get("session_id") + if session_id is not None: + return str(session_id) + + return None + + def _get_max_iterations( + self, user_api_key_dict: UserAPIKeyAuth + ) -> Optional[int]: + """Extract max_iterations from key metadata.""" + metadata = user_api_key_dict.metadata or {} + max_iterations = metadata.get("max_iterations") + if max_iterations is not None: + return int(max_iterations) + return None + + def _make_cache_key(self, session_id: str) -> str: + """ + Create cache key for session iteration counter. + + Uses Redis hash-tag pattern {session_iterations:} so all + keys for a session land on the same Redis Cluster slot. + """ + return f"{{session_iterations:{session_id}}}:count" + + async def _increment_and_get(self, cache_key: str) -> int: + """ + Atomically increment the session counter and return the new value. + + Tries Redis first (via registered Lua script for atomicity across + instances), falls back to in-memory cache. + """ + if self.increment_script is not None: + try: + result = await self.increment_script( + keys=[cache_key], + args=[self.ttl], + ) + return int(result) + except Exception as e: + verbose_proxy_logger.warning( + "MaxIterationsHandler: Redis failed, falling back to in-memory: %s", + str(e), + ) + + # Fallback: in-memory cache + return await self._in_memory_increment(cache_key) + + async def _in_memory_increment(self, cache_key: str) -> int: + """Increment counter in in-memory cache with TTL.""" + current = await self.internal_usage_cache.async_get_cache( + key=cache_key, + litellm_parent_otel_span=None, + local_only=True, + ) + new_value = (int(current) if current is not None else 0) + 1 + await self.internal_usage_cache.async_set_cache( + key=cache_key, + value=new_value, + ttl=self.ttl, + litellm_parent_otel_span=None, + local_only=True, + ) + return new_value diff --git a/litellm/types/agents.py b/litellm/types/agents.py index f4e410a3e2d..3ad898b1935 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -167,16 +167,25 @@ class AugmentedAgentCard(AgentCard): is_public: bool +# Object permission shape for agent MCP tool access (mirrors LiteLLM_ObjectPermissionBase) +class AgentObjectPermission(TypedDict, total=False): + mcp_servers: Optional[List[str]] + mcp_access_groups: Optional[List[str]] + mcp_tool_permissions: Optional[Dict[str, List[str]]] + + class AgentConfig(TypedDict, total=False): agent_name: Required[str] agent_card_params: Required[AgentCard] litellm_params: Dict[str, Any] # allow for any future litellm params + object_permission: AgentObjectPermission class PatchAgentRequest(TypedDict, total=False): agent_name: str agent_card_params: AgentCard litellm_params: Dict[str, Any] + object_permission: AgentObjectPermission # Request/Response models for CRUD endpoints @@ -187,6 +196,7 @@ class AgentResponse(BaseModel): agent_name: str litellm_params: Optional[Dict[str, Any]] = None agent_card_params: Dict[str, Any] + object_permission: Optional[Dict[str, Any]] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None created_by: Optional[str] = None diff --git a/schema.prisma b/schema.prisma index 4af7484148c..155cea12ca4 100644 --- a/schema.prisma +++ b/schema.prisma @@ -64,6 +64,8 @@ model LiteLLM_AgentsTable { litellm_params Json? agent_card_params Json agent_access_groups String[] @default([]) + object_permission_id String? + object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@ -264,6 +266,7 @@ model LiteLLM_ObjectPermissionTable { organizations LiteLLM_OrganizationTable[] users LiteLLM_UserTable[] end_users LiteLLM_EndUserTable[] + agents_table LiteLLM_AgentsTable[] } // Holds the MCP server configuration @@ -273,7 +276,6 @@ model LiteLLM_MCPServerTable { alias String? description String? url String? - spec_path String? transport String @default("sse") auth_type String? credentials Json? @default("{}") @@ -315,6 +317,7 @@ model LiteLLM_VerificationToken { router_settings Json? @default("{}") user_id String? team_id String? + agent_id String? project_id String? permissions Json @default("{}") max_parallel_requests Int? @@ -1052,6 +1055,26 @@ model LiteLLM_PolicyAttachmentTable { updated_by String? } +// Global tool registry - auto-discovered from LLM responses; admins set call_policy here +model LiteLLM_ToolTable { + tool_id String @id @default(uuid()) + tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space" + origin String? // MCP server name or "user_defined" + call_policy String @default("untrusted") // "trusted" | "untrusted" | "dual_llm" | "blocked" + call_count Int @default(0) // cumulative number of times this tool was seen + assignments Json? @default("{}") + key_hash String? // hash of the virtual key that first called this tool + team_id String? // team that first called this tool + key_alias String? // human-readable alias of the virtual key + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? + + @@index([call_policy]) + @@index([team_id]) +} + //Unified Access Groups table for storing unified access groups model LiteLLM_AccessGroupTable { access_group_id String @id @default(uuid()) diff --git a/scripts/test_agent_mcp_endpoints.sh b/scripts/test_agent_mcp_endpoints.sh new file mode 100755 index 00000000000..93cc68db2ed --- /dev/null +++ b/scripts/test_agent_mcp_endpoints.sh @@ -0,0 +1,186 @@ +#!/usr/bin/env bash +# +# Test agent endpoint-level changes for MCP tool permissions (object_permission). +# Requires: proxy running, valid admin API key, curl, jq. +# +# Usage: +# export LITELLM_PROXY_BASE_URL="http://localhost:4000" # optional, default below +# export LITELLM_API_KEY="sk-..." # required +# ./scripts/test_agent_mcp_endpoints.sh +# +set -euo pipefail + +BASE_URL="${LITELLM_PROXY_BASE_URL:-http://localhost:4000}" +API_KEY="${LITELLM_API_KEY:-}" + +if ! command -v jq &>/dev/null; then + echo "Error: jq is required. Install with: brew install jq (macOS) or apt install jq (Linux)" + exit 1 +fi +if [[ -z "$API_KEY" ]]; then + echo "Error: LITELLM_API_KEY is not set. Export it or pass via env." + exit 1 +fi + +AUTH_HEADER="Authorization: Bearer $API_KEY" +AGENT_NAME="test-agent-mcp-$(date +%s)" + +# Colors for output +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +NC='\033[0m' + +pass() { echo -e "${GREEN}PASS${NC}: $*"; } +fail() { echo -e "${RED}FAIL${NC}: $*"; exit 1; } +info() { echo -e "${YELLOW}INFO${NC}: $*"; } + +# --- 1. Create agent with object_permission --- +info "Creating agent with object_permission (mcp_servers, mcp_tool_permissions)..." +CREATE_RESP=$(curl -s -w "\n%{http_code}" -X POST "$BASE_URL/v1/agents" \ + -H "$AUTH_HEADER" \ + -H "Content-Type: application/json" \ + -d '{ + "agent_name": "'"$AGENT_NAME"'", + "agent_card_params": { + "protocolVersion": "1.0", + "name": "Test MCP Agent", + "description": "Agent for endpoint tests", + "url": "http://localhost:9999/", + "version": "1.0.0", + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "capabilities": {"streaming": true}, + "skills": [] + }, + "object_permission": { + "mcp_servers": ["server_1", "server_2"], + "mcp_access_groups": ["group_a"], + "mcp_tool_permissions": {"server_1": ["tool_a", "tool_b"], "server_2": ["tool_c"]} + } + }') +HTTP_CODE=$(echo "$CREATE_RESP" | tail -n1) +BODY=$(echo "$CREATE_RESP" | sed '$d') +if [[ "$HTTP_CODE" != "200" ]]; then + fail "POST /v1/agents returned $HTTP_CODE. Body: $BODY" +fi +AGENT_ID=$(echo "$BODY" | jq -r '.agent_id') +if [[ -z "$AGENT_ID" || "$AGENT_ID" == "null" ]]; then + fail "POST /v1/agents did not return agent_id. Body: $BODY" +fi +pass "Created agent $AGENT_ID" + +# Check create response includes object_permission +OP=$(echo "$BODY" | jq '.object_permission') +if [[ "$OP" == "null" || -z "$OP" ]]; then + fail "POST /v1/agents response missing object_permission. Body: $BODY" +fi +SERVERS=$(echo "$OP" | jq -r '.mcp_servers | join(",")') +if [[ "$SERVERS" != "server_1,server_2" ]]; then + fail "object_permission.mcp_servers unexpected: $SERVERS" +fi +pass "Create response includes object_permission with mcp_servers and mcp_tool_permissions" + +# --- 2. GET /v1/agents (list) includes object_permission for our agent --- +info "GET /v1/agents and check one agent has object_permission..." +LIST_RESP=$(curl -s -w "\n%{http_code}" -X GET "$BASE_URL/v1/agents" -H "$AUTH_HEADER") +LIST_CODE=$(echo "$LIST_RESP" | tail -n1) +LIST_BODY=$(echo "$LIST_RESP" | sed '$d') +if [[ "$LIST_CODE" != "200" ]]; then + fail "GET /v1/agents returned $LIST_CODE" +fi +AGENT_IN_LIST=$(echo "$LIST_BODY" | jq --arg id "$AGENT_ID" '.[] | select(.agent_id == $id)') +if [[ -z "$AGENT_IN_LIST" ]]; then + fail "GET /v1/agents did not return agent $AGENT_ID (list might be key-scoped)" +fi +OP_LIST=$(echo "$AGENT_IN_LIST" | jq '.object_permission') +if [[ "$OP_LIST" == "null" || -z "$OP_LIST" ]]; then + fail "GET /v1/agents list entry for agent missing object_permission" +fi +pass "GET /v1/agents list includes object_permission for agent" + +# --- 3. GET /v1/agents/{agent_id} returns object_permission --- +info "GET /v1/agents/{agent_id}..." +GET_RESP=$(curl -s -w "\n%{http_code}" -X GET "$BASE_URL/v1/agents/$AGENT_ID" -H "$AUTH_HEADER") +GET_CODE=$(echo "$GET_RESP" | tail -n1) +GET_BODY=$(echo "$GET_RESP" | sed '$d') +if [[ "$GET_CODE" != "200" ]]; then + fail "GET /v1/agents/$AGENT_ID returned $GET_CODE. Body: $GET_BODY" +fi +OP_GET=$(echo "$GET_BODY" | jq '.object_permission') +if [[ "$OP_GET" == "null" || -z "$OP_GET" ]]; then + fail "GET /v1/agents/$AGENT_ID response missing object_permission" +fi +TOOL_PERMS=$(echo "$OP_GET" | jq -r '.mcp_tool_permissions.server_1 | join(",")') +if [[ "$TOOL_PERMS" != "tool_a,tool_b" ]]; then + fail "object_permission.mcp_tool_permissions.server_1 unexpected: $TOOL_PERMS" +fi +pass "GET /v1/agents/{agent_id} returns object_permission with mcp_tool_permissions" + +# --- 4. PATCH /v1/agents/{agent_id} with new object_permission --- +info "PATCH /v1/agents/{agent_id} with updated object_permission..." +PATCH_RESP=$(curl -s -w "\n%{http_code}" -X PATCH "$BASE_URL/v1/agents/$AGENT_ID" \ + -H "$AUTH_HEADER" \ + -H "Content-Type: application/json" \ + -d '{ + "object_permission": { + "mcp_servers": ["server_3"], + "mcp_tool_permissions": {"server_3": ["tool_x"]} + } + }') +PATCH_CODE=$(echo "$PATCH_RESP" | tail -n1) +PATCH_BODY=$(echo "$PATCH_RESP" | sed '$d') +if [[ "$PATCH_CODE" != "200" ]]; then + fail "PATCH /v1/agents/$AGENT_ID returned $PATCH_CODE. Body: $PATCH_BODY" +fi +OP_PATCH=$(echo "$PATCH_BODY" | jq '.object_permission') +if [[ "$OP_PATCH" == "null" || -z "$OP_PATCH" ]]; then + fail "PATCH response missing object_permission" +fi +PATCH_SERVERS=$(echo "$OP_PATCH" | jq -r '.mcp_servers | join(",")') +if [[ "$PATCH_SERVERS" != "server_3" ]]; then + fail "PATCH object_permission.mcp_servers unexpected: $PATCH_SERVERS" +fi +pass "PATCH /v1/agents/{agent_id} updates and returns object_permission" + +# --- 5. Create agent without object_permission; GET should still work --- +info "Creating agent without object_permission..." +AGENT_NAME_2="test-agent-no-mcp-$(date +%s)" +CREATE2_RESP=$(curl -s -w "\n%{http_code}" -X POST "$BASE_URL/v1/agents" \ + -H "$AUTH_HEADER" \ + -H "Content-Type: application/json" \ + -d '{ + "agent_name": "'"$AGENT_NAME_2"'", + "agent_card_params": { + "protocolVersion": "1.0", + "name": "No MCP Agent", + "description": "No object_permission", + "url": "http://localhost:9999/", + "version": "1.0.0", + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "capabilities": {}, + "skills": [] + } + }') +CODE2=$(echo "$CREATE2_RESP" | tail -n1) +BODY2=$(echo "$CREATE2_RESP" | sed '$d') +if [[ "$CODE2" != "200" ]]; then + fail "POST /v1/agents (no object_permission) returned $CODE2. Body: $BODY2" +fi +AGENT_ID_2=$(echo "$BODY2" | jq -r '.agent_id') +# object_permission may be null or absent +pass "Created agent without object_permission: $AGENT_ID_2" + +# --- 6. Cleanup: delete both agents --- +info "Deleting test agents..." +for AID in "$AGENT_ID" "$AGENT_ID_2"; do + DEL_CODE=$(curl -s -o /dev/null -w "%{http_code}" -X DELETE "$BASE_URL/v1/agents/$AID" -H "$AUTH_HEADER") + if [[ "$DEL_CODE" != "200" ]]; then + info "DELETE /v1/agents/$AID returned $DEL_CODE (non-fatal)" + fi +done +pass "Cleanup done" + +echo "" +echo -e "${GREEN}All endpoint checks passed.${NC}" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index c2dbc94f721..b7ae33d1f80 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -1595,3 +1595,146 @@ async def test_get_allowed_mcp_servers_for_key_prefers_in_memory_permission(): assert set(result) == {"direct-server", "group-server"} mock_get_perm.assert_not_called() mock_access_groups.assert_called_once_with(["grp-alpha"]) + + +@pytest.mark.asyncio +class TestAgentMCPPermissions: + """Test agent-level MCP server and tool permission intersection.""" + + async def test_get_allowed_mcp_servers_agent_intersection(self): + """Key/team allow [server_1, server_2]; agent allows [server_1]. Result = [server_1].""" + user_api_key_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + agent_id="agent-123", + ) + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_key" + ) as mock_key: + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_team" + ) as mock_team: + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_agent" + ) as mock_agent: + mock_key.return_value = ["server_1", "server_2"] + mock_team.return_value = [] + mock_agent.return_value = ["server_1"] + result = await MCPRequestHandler.get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth + ) + assert sorted(result) == ["server_1"] + mock_agent.assert_called_once_with(user_api_key_auth) + + async def test_get_allowed_mcp_servers_agent_no_restriction(self): + """Agent with no object_permission returns []; no intersection applied (inherit key/team).""" + user_api_key_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + agent_id="agent-456", + ) + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_key" + ) as mock_key: + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_team" + ) as mock_team: + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_agent" + ) as mock_agent: + mock_key.return_value = ["server_1", "server_2"] + mock_team.return_value = [] + mock_agent.return_value = [] # no agent-level restriction + result = await MCPRequestHandler.get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth + ) + assert sorted(result) == ["server_1", "server_2"] + mock_agent.assert_called_once_with(user_api_key_auth) + + async def test_get_allowed_mcp_servers_key_team_agent_intersection(self): + """Key allows [1, 2], agent allows [2, 3]. Result = [2].""" + user_api_key_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + agent_id="agent-789", + ) + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_key" + ) as mock_key: + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_team" + ) as mock_team: + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_agent" + ) as mock_agent: + mock_key.return_value = ["server_1", "server_2"] + mock_team.return_value = [] + mock_agent.return_value = ["server_2", "server_3"] + result = await MCPRequestHandler.get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth + ) + assert sorted(result) == ["server_2"] + + async def test_get_allowed_tools_for_server_agent_intersection(self): + """Key allows [tool_a, tool_b], agent allows [tool_a]. Result = [tool_a].""" + user_api_key_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + agent_id="agent-tools", + ) + key_perm = MagicMock() + key_perm.mcp_tool_permissions = {"server_1": ["tool_a", "tool_b"]} + team_perm = None + with patch.object( + MCPRequestHandler, "_get_key_object_permission", return_value=key_perm + ): + with patch.object( + MCPRequestHandler, "_get_team_object_permission", + new_callable=AsyncMock, + return_value=team_perm, + ): + with patch.object( + MCPRequestHandler, + "_get_agent_tool_permissions_for_server", + new_callable=AsyncMock, + return_value=["tool_a"], + ) as mock_agent_tools: + result = await MCPRequestHandler.get_allowed_tools_for_server( + server_id="server_1", + user_api_key_auth=user_api_key_auth, + ) + assert result == ["tool_a"] + mock_agent_tools.assert_called_once() + call_kwargs = mock_agent_tools.call_args.kwargs + assert call_kwargs["server_id"] == "server_1" + assert call_kwargs["user_api_key_auth"] == user_api_key_auth + + async def test_get_allowed_tools_for_server_agent_no_restriction(self): + """Agent has no tool permissions for server; key/team result is unchanged.""" + user_api_key_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + agent_id="agent-no-tools", + ) + key_perm = MagicMock() + key_perm.mcp_tool_permissions = {"server_1": ["tool_a", "tool_b"]} + with patch.object( + MCPRequestHandler, "_get_key_object_permission", return_value=key_perm + ): + with patch.object( + MCPRequestHandler, "_get_team_object_permission", + new_callable=AsyncMock, + return_value=None, + ): + with patch.object( + MCPRequestHandler, + "_get_agent_tool_permissions_for_server", + new_callable=AsyncMock, + return_value=None, + ): + result = await MCPRequestHandler.get_allowed_tools_for_server( + server_id="server_1", + user_api_key_auth=user_api_key_auth, + ) + assert sorted(result) == ["tool_a", "tool_b"] diff --git a/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py b/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py new file mode 100644 index 00000000000..deb1c483b87 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py @@ -0,0 +1,106 @@ +""" +Unit Tests for the max iterations limiter for the proxy. + +Tests that session-scoped iteration counting works correctly: +- Enforces max_iterations per session_id +- Different sessions have independent counters +""" + +import pytest +from fastapi import HTTPException + +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.max_iterations_limiter import _PROXY_MaxIterationsHandler +from litellm.proxy.utils import InternalUsageCache + + +@pytest.mark.asyncio +async def test_max_iterations_basic_enforcement(): + """ + Test that max_iterations is enforced per session_id. + + - 3 requests with the same session_id should succeed when max_iterations=3 + - 4th request should raise 429 + """ + local_cache = DualCache() + handler = _PROXY_MaxIterationsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-1234", metadata={"max_iterations": 3} + ) + + # First 3 requests should succeed + for i in range(3): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-abc"}}, + call_type="", + ) + + # 4th request should fail with 429 + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-abc"}}, + call_type="", + ) + assert exc_info.value.status_code == 429 + assert "max_iterations" in str(exc_info.value.detail).lower() + + +@pytest.mark.asyncio +async def test_max_iterations_different_sessions_independent(): + """ + Test that different session_ids have independent iteration counters. + + - Session A and Session B each get their own max_iterations budget + - Exhausting Session A does not affect Session B + """ + local_cache = DualCache() + handler = _PROXY_MaxIterationsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-5678", metadata={"max_iterations": 2} + ) + + # Session A: 2 calls succeed + for _ in range(2): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-A"}}, + call_type="", + ) + + # Session B: 2 calls succeed (independent counter) + for _ in range(2): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-B"}}, + call_type="", + ) + + # Session A: 3rd call fails + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-A"}}, + call_type="", + ) + assert exc_info.value.status_code == 429 + + # Session B: 3rd call also fails + with pytest.raises(HTTPException): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-B"}}, + call_type="", + ) diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index fc2aa1599d3..cc04e674003 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -13056,6 +13056,21 @@ "type": "github", "url": "https://github.com/sponsors/wooorm" } + }, + "node_modules/@next/swc-win32-ia32-msvc": { + "version": "14.2.33", + "resolved": "https://registry.npmjs.org/@next/swc-win32-ia32-msvc/-/swc-win32-ia32-msvc-14.2.33.tgz", + "integrity": "sha512-pc9LpGNKhJ0dXQhZ5QMmYxtARwwmWLpeocFmVG5Z0DzWq5Uf0izcI8tLc+qOpqxO1PWqZ5A7J1blrUIKrIFc7Q==", + "cpu": [ + "ia32" + ], + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">= 10" + } } } } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelRetrySettingsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelRetrySettingsTab.test.tsx index 32e619d084c..5b756a833d8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelRetrySettingsTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelRetrySettingsTab.test.tsx @@ -1,7 +1,7 @@ import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import React from "react"; -import { describe, beforeEach, expect, it, vi } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import ModelRetrySettingsTab from "./ModelRetrySettingsTab"; // TabPanel requires a parent Tabs context in Tremor. We stub it to render children diff --git a/ui/litellm-dashboard/src/components/agents.tsx b/ui/litellm-dashboard/src/components/agents.tsx index d6d8b320ec9..7ab6a5f0381 100644 --- a/ui/litellm-dashboard/src/components/agents.tsx +++ b/ui/litellm-dashboard/src/components/agents.tsx @@ -1,13 +1,13 @@ import React, { useState, useEffect } from "react"; import { Button } from "@tremor/react"; -import { Modal } from "antd"; -import { getAgentsList, deleteAgentCall } from "./networking"; +import { Modal, Alert } from "antd"; +import { getAgentsList, deleteAgentCall, keyListCall } from "./networking"; import AddAgentForm from "./agents/add_agent_form"; -import AgentTable from "./agents/agent_table"; +import AgentCardGrid from "./agents/agent_card_grid"; import { isAdminRole } from "@/utils/roles"; import AgentInfoView from "./agents/agent_info"; import NotificationsManager from "./molecules/notifications_manager"; -import { Agent } from "./agents/types"; +import { Agent, AgentKeyInfo } from "./agents/types"; interface AgentsPanelProps { accessToken: string | null; @@ -20,6 +20,7 @@ interface AgentsResponse { const AgentsPanel: React.FC = ({ accessToken, userRole }) => { const [agentsList, setAgentsList] = useState([]); + const [keyInfoMap, setKeyInfoMap] = useState>({}); const [isAddModalVisible, setIsAddModalVisible] = useState(false); const [isLoading, setIsLoading] = useState(false); const [isDeleting, setIsDeleting] = useState(false); @@ -36,8 +37,7 @@ const AgentsPanel: React.FC = ({ accessToken, userRole }) => { setIsLoading(true); try { const response: AgentsResponse = await getAgentsList(accessToken); - console.log(`agents: ${JSON.stringify(response)}`); - setAgentsList(response.agents); + setAgentsList(response.agents || []); } catch (error) { console.error("Error fetching agents:", error); } finally { @@ -45,10 +45,50 @@ const AgentsPanel: React.FC = ({ accessToken, userRole }) => { } }; + const fetchKeysForAgents = async () => { + if (!accessToken) return; + try { + const { keys = [] } = await keyListCall( + accessToken, + null, + null, + null, + null, + null, + 1, + 500 + ); + const map: Record = {}; + for (const key of keys) { + const agentId = (key as { agent_id?: string }).agent_id; + if (agentId && !map[agentId]) { + map[agentId] = { + has_key: true, + key_alias: (key as { key_alias?: string }).key_alias, + token_prefix: (key as { token?: string }).token + ? `${(key as { token: string }).token.slice(0, 8)}…` + : undefined, + }; + } + } + setKeyInfoMap(map); + } catch (error) { + console.error("Error fetching keys for agents:", error); + } + }; + useEffect(() => { fetchAgents(); }, [accessToken]); + useEffect(() => { + if (accessToken && agentsList.length > 0) { + fetchKeysForAgents(); + } else if (agentsList.length === 0) { + setKeyInfoMap({}); + } + }, [accessToken, agentsList.length]); + const handleAddAgent = () => { if (selectedAgentId) { setSelectedAgentId(null); @@ -94,6 +134,13 @@ const AgentsPanel: React.FC = ({ accessToken, userRole }) => {

Agents

List of A2A-spec agents that are available to be used in your organization. Go to AI Hub, to make agents public.

+