From 0b86fc56f2d6d7d62e1b6c6e9bc5566feaa07c06 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 24 Feb 2026 20:56:07 -0800 Subject: [PATCH] feat(agent_endpoints/): allow giving agents keys --- .../migration.sql | 40 ++++ .../litellm_proxy_extras/schema.prisma | 25 ++- .../mcp_server/auth/user_api_key_auth_mcp.py | 190 ++++++++++++++---- litellm/proxy/_types.py | 76 +++---- .../proxy/agent_endpoints/agent_registry.py | 107 ++++++++-- litellm/proxy/agent_endpoints/endpoints.py | 58 +++--- litellm/types/agents.py | 10 + schema.prisma | 25 ++- scripts/test_agent_mcp_endpoints.sh | 186 +++++++++++++++++ .../auth/test_user_api_key_auth_mcp.py | 156 +++++++++++++- .../components/ModelRetrySettingsTab.test.tsx | 2 +- .../src/components/ToolPolicies.tsx | 5 +- .../src/components/agents.tsx | 62 +++++- .../src/components/agents/add_agent_form.tsx | 103 ++++++++-- .../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 | 3 +- .../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 +- 30 files changed, 1114 insertions(+), 181 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260224203854_add_agent_object_permissions_table/migration.sql create mode 100755 scripts/test_agent_mcp_endpoints.sh 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..dd987372df8 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 @@ -6,12 +6,8 @@ from starlette.requests import Request from starlette.types import Scope from litellm._logging import verbose_logger -from litellm.proxy._types import ( - LiteLLM_TeamTable, - ProxyException, - SpecialHeaders, - UserAPIKeyAuth, -) +from litellm.proxy._types import (LiteLLM_TeamTable, ProxyException, + SpecialHeaders, UserAPIKeyAuth) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -412,6 +408,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)}") @@ -443,11 +459,9 @@ class MCPRequestHandler: get_team_object() in litellm/proxy/auth/auth_checks.py """ from litellm.proxy.auth.auth_checks import get_team_object - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + from litellm.proxy.proxy_server import (prisma_client, + proxy_logging_obj, + user_api_key_cache) verbose_logger.debug( f"MCP team permission lookup: team_id={user_api_key_auth.team_id if user_api_key_auth else None}" @@ -513,13 +527,28 @@ 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: + agent_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( + server_id=server_id, + user_api_key_auth=user_api_key_auth, + ) + 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)}") @@ -582,12 +611,11 @@ class MCPRequestHandler: user_api_key_auth ) if key_object_permission is None and user_api_key_auth and user_api_key_auth.object_permission_id: - from litellm.proxy.auth.auth_checks import get_object_permission - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + from litellm.proxy.auth.auth_checks import \ + get_object_permission + from litellm.proxy.proxy_server import (prisma_client, + proxy_logging_obj, + user_api_key_cache) if prisma_client is not None: key_object_permission = await get_object_permission( object_permission_id=user_api_key_auth.object_permission_id, @@ -665,11 +693,9 @@ class MCPRequestHandler: Returns the MCP servers from the end_user's object_permission. """ from litellm.proxy.auth.auth_checks import get_end_user_object - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + from litellm.proxy.proxy_server import (prisma_client, + proxy_logging_obj, + user_api_key_cache) if not user_api_key_auth or not user_api_key_auth.end_user_id: return [] @@ -715,6 +741,99 @@ class MCPRequestHandler: ) return [] + @staticmethod + async def _get_allowed_mcp_servers_for_agent( + user_api_key_auth: Optional[UserAPIKeyAuth] = 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). + """ + from litellm.proxy.proxy_server import prisma_client + + if not user_api_key_auth or not user_api_key_auth.agent_id: + return [] + + if prisma_client is None: + verbose_logger.debug("prisma_client is None") + return [] + + 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 [] + + obj_perm = agent_row.object_permission + 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, + ) -> 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. + """ + from litellm.proxy.proxy_server import prisma_client + + if not user_api_key_auth or not user_api_key_auth.agent_id or not prisma_client: + 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 + + obj_perm = agent_row.object_permission + 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] @@ -761,9 +880,8 @@ class MCPRequestHandler: try: # Import here to avoid circular import - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import \ + global_mcp_server_manager # Use the new helper for config-loaded servers server_ids = MCPRequestHandler._get_config_server_ids_for_access_groups( @@ -817,11 +935,9 @@ class MCPRequestHandler: user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> List[str]: from litellm.proxy.auth.auth_checks import get_object_permission - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + from litellm.proxy.proxy_server import (prisma_client, + proxy_logging_obj, + user_api_key_cache) if user_api_key_auth is None: return [] @@ -857,11 +973,9 @@ class MCPRequestHandler: Get MCP access groups for the team """ from litellm.proxy.auth.auth_checks import get_team_object - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + from litellm.proxy.proxy_server import (prisma_client, + proxy_logging_obj, + user_api_key_cache) if user_api_key_auth is None: return [] 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..4feddd67dd7 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -5,6 +5,8 @@ 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 ( + attach_object_permission_to_dict, handle_update_object_permission_common) from litellm.proxy.utils import PrismaClient from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -117,20 +119,36 @@ 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 ) - return AgentResponse(**created_agent.model_dump()) # type: ignore + created_agent_dict = created_agent.model_dump() + await attach_object_permission_to_dict( + created_agent_dict, prisma_client + ) + return AgentResponse(**created_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error adding agent to DB: {str(e)}") @@ -181,7 +199,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 +210,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}, @@ -201,7 +233,11 @@ class AgentRegistry: "updated_at": datetime.now(timezone.utc), }, ) - return AgentResponse(**patched_agent.model_dump()) # type: ignore + patched_agent_dict = patched_agent.model_dump() + await attach_object_permission_to_dict( + patched_agent_dict, prisma_client + ) + return AgentResponse(**patched_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error patching agent in DB: {str(e)}") @@ -238,19 +274,44 @@ 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, ) - return AgentResponse(**updated_agent.model_dump()) # type: ignore + updated_agent_dict = updated_agent.model_dump() + await attach_object_permission_to_dict( + updated_agent_dict, prisma_client + ) + return AgentResponse(**updated_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error updating agent in DB: {str(e)}") @@ -268,7 +329,11 @@ class AgentRegistry: agents: List[Dict[str, Any]] = [] for agent in agents_from_db: - agents.append(dict(agent)) + agent_dict = dict(agent) + await attach_object_permission_to_dict( + agent_dict, prisma_client + ) + 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..5200a1fc94d 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -14,20 +14,16 @@ from fastapi import APIRouter, Depends, HTTPException, Request import litellm from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import (CommonProxyErrors, LitellmUserRoles, + UserAPIKeyAuth) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.types.agents import ( - AgentConfig, - AgentMakePublicResponse, - AgentResponse, - 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, -) +from litellm.proxy.management_endpoints.common_daily_activity import \ + get_daily_activity +from litellm.types.agents import (AgentConfig, AgentMakePublicResponse, + AgentResponse, MakeAgentsPublicRequest, + PatchAgentRequest) +from litellm.types.proxy.management_endpoints.common_daily_activity import \ + SpendAnalyticsPaginatedResponse router = APIRouter() @@ -53,10 +49,10 @@ async def get_agents( Returns: List[AgentResponse] """ - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( - AgentRequestHandler, - ) + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import \ + AgentRequestHandler try: returned_agents: List[AgentResponse] = [] @@ -109,9 +105,8 @@ async def get_agents( #### CRUD ENDPOINTS FOR AGENTS #### -from litellm.proxy.agent_endpoints.agent_registry import ( - global_agent_registry as AGENT_REGISTRY, -) +from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry as AGENT_REGISTRY @router.post( @@ -231,13 +226,20 @@ async def get_agent_by_id(agent_id: str): raise HTTPException(status_code=500, detail="Prisma client not initialized") try: + from litellm.proxy.management_helpers.object_permission_utils import \ + attach_object_permission_to_dict + agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: - agent = await prisma_client.db.litellm_agentstable.find_unique( + agent_row = await prisma_client.db.litellm_agentstable.find_unique( where={"agent_id": agent_id} ) - if agent is not None: - agent = AgentResponse(**agent.model_dump()) # type: ignore + if agent_row is not None: + agent_dict = agent_row.model_dump() + await attach_object_permission_to_dict( + agent_dict, prisma_client + ) + agent = AgentResponse(**agent_dict) # type: ignore if agent is None: raise HTTPException( @@ -530,9 +532,8 @@ async def make_agent_public( try: # Update the public model groups import litellm - from litellm.proxy.agent_endpoints.agent_registry import ( - global_agent_registry as AGENT_REGISTRY, - ) + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry as AGENT_REGISTRY from litellm.proxy.proxy_server import proxy_config # Check if user has admin permissions @@ -647,9 +648,8 @@ async def make_agents_public( try: # Update the public model groups import litellm - from litellm.proxy.agent_endpoints.agent_registry import ( - global_agent_registry as AGENT_REGISTRY, - ) + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry as AGENT_REGISTRY from litellm.proxy.proxy_server import proxy_config # Load existing config 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..6232f7974da 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 @@ -15,9 +15,8 @@ sys.path.insert( from starlette.datastructures import Headers -from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( - MCPRequestHandler, -) +from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import \ + MCPRequestHandler from litellm.proxy._types import SpecialHeaders, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -1176,7 +1175,8 @@ class TestMCPAccessGroupsE2E: @pytest.mark.asyncio def test_mcp_path_based_server_segregation(monkeypatch): # Import the MCP server FastAPI app and context getter - from litellm.proxy._experimental.mcp_server.server import app, get_auth_context + from litellm.proxy._experimental.mcp_server.server import ( + app, get_auth_context) captured_mcp_servers = {} @@ -1277,7 +1277,8 @@ async def test_get_team_object_permission_with_already_loaded_permission(): Test that _get_team_object_permission returns the already loaded object_permission from the team object without making an additional DB call. """ - from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable + from litellm.proxy._types import (LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTable) # Create mock object permission mock_object_permission = LiteLLM_ObjectPermissionTable( @@ -1340,7 +1341,8 @@ async def test_get_team_object_permission_with_core_auth_auto_loading(): the team object returned by get_team_object() should already have object_permission loaded when an object_permission_id exists. """ - from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable + from litellm.proxy._types import (LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTable) # Create mock object permission mock_object_permission = LiteLLM_ObjectPermissionTable( @@ -1595,3 +1597,145 @@ 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_with( + "server_1", 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/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/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx index 0e3f5434e7f..1e1a9b3b13d 100644 --- a/ui/litellm-dashboard/src/components/ToolPolicies.tsx +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -2,7 +2,6 @@ import React, { useCallback, useDeferredValue, useEffect, useState } from "react"; import { Select, Switch, Tooltip } from "antd"; -import { Select, Tooltip } from "antd"; import { Table, TableHead, @@ -68,7 +67,7 @@ const PolicySelect: React.FC<{ paddingLeft: 8, paddingRight: 4, }, - }} + } as Record} popupMatchSelectWidth={false} options={POLICY_OPTIONS.map((o) => ({ value: o.value, @@ -266,7 +265,7 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
Live Tail - +
@@ -577,11 +651,16 @@ const AddAgentForm: React.FC = ({ )} {currentStep === 1 && ( + + )} + {currentStep === 2 && ( )} - {currentStep === 2 && ( + {currentStep === 3 && ( diff --git a/ui/litellm-dashboard/src/components/agents/agent_card.tsx b/ui/litellm-dashboard/src/components/agents/agent_card.tsx new file mode 100644 index 00000000000..dc684386edf --- /dev/null +++ b/ui/litellm-dashboard/src/components/agents/agent_card.tsx @@ -0,0 +1,103 @@ +import React from "react"; +import { Card, Badge, Tooltip, Button } from "antd"; +import { CopyOutlined, KeyOutlined, WarningOutlined, DeleteOutlined } from "@ant-design/icons"; +import { Agent, AgentKeyInfo } from "./types"; + +interface AgentCardProps { + agent: Agent; + keyInfo?: AgentKeyInfo; + onAgentClick: (agentId: string) => void; + onDeleteClick?: (agentId: string, agentName: string) => void; + accessToken: string | null; + isAdmin: boolean; + onAgentUpdated: () => void; +} + +const AgentCard: React.FC = ({ + agent, + keyInfo, + onAgentClick, + onDeleteClick, + isAdmin, +}) => { + const description = + agent.agent_card_params?.description || "No description"; + const url = agent.agent_card_params?.url; + const hasKey = keyInfo?.has_key ?? false; + const statusBadge = hasKey ? ( + + ) : ( + + ); + + const copyToClipboard = (e: React.MouseEvent, text: string) => { + e.stopPropagation(); + navigator.clipboard.writeText(text); + }; + + return ( + onAgentClick(agent.agent_id)} + > +
+
+
+ + {agent.agent_name} + + + copyToClipboard(e, agent.agent_id)} + className="cursor-pointer text-gray-400 hover:text-blue-500 text-xs shrink-0" + /> + +
+
{statusBadge}
+
+ {isAdmin && onDeleteClick && ( + +
+

+ {description} +

+ {url && ( +

+ {url} +

+ )} +
+ {hasKey ? ( +
+ + {keyInfo?.key_alias || keyInfo?.token_prefix || "Key assigned"} +
+ ) : ( +
+ + No key assigned +
+ )} +
+ + ); +}; + +export default AgentCard; diff --git a/ui/litellm-dashboard/src/components/agents/agent_card_grid.tsx b/ui/litellm-dashboard/src/components/agents/agent_card_grid.tsx new file mode 100644 index 00000000000..0ba7902f8b2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/agents/agent_card_grid.tsx @@ -0,0 +1,63 @@ +import React from "react"; +import { Skeleton } from "antd"; +import AgentCard from "./agent_card"; +import { Agent, AgentKeyInfo } from "./types"; + +interface AgentCardGridProps { + agentsList: Agent[]; + keyInfoMap: Record; + isLoading: boolean; + onDeleteClick: (agentId: string, agentName: string) => void; + accessToken: string | null; + onAgentUpdated: () => void; + isAdmin: boolean; + onAgentClick: (agentId: string) => void; +} + +const AgentCardGrid: React.FC = ({ + agentsList, + keyInfoMap, + isLoading, + onDeleteClick, + accessToken, + onAgentUpdated, + isAdmin, + onAgentClick, +}) => { + if (isLoading) { + return ( +
+ {[1, 2, 3].map((i) => ( + + ))} +
+ ); + } + + if (!agentsList || agentsList.length === 0) { + return ( +
+

No agents found. Create one to get started.

+
+ ); + } + + return ( +
+ {agentsList.map((agent) => ( + + ))} +
+ ); +}; + +export default AgentCardGrid; diff --git a/ui/litellm-dashboard/src/components/agents/agent_info.tsx b/ui/litellm-dashboard/src/components/agents/agent_info.tsx index 8d0febd8417..deb4f900377 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_info.tsx +++ b/ui/litellm-dashboard/src/components/agents/agent_info.tsx @@ -205,6 +205,44 @@ const AgentInfoView: React.FC = ({ {formatDate(agent.updated_at)} + {agent.object_permission && + (agent.object_permission.mcp_servers?.length || + agent.object_permission.mcp_access_groups?.length || + (agent.object_permission.mcp_tool_permissions && + Object.keys(agent.object_permission.mcp_tool_permissions).length > 0)) && ( +
+ MCP Tool Permissions + + {agent.object_permission.mcp_servers && agent.object_permission.mcp_servers.length > 0 && ( + + {agent.object_permission.mcp_servers.join(", ")} + + )} + {agent.object_permission.mcp_access_groups && + agent.object_permission.mcp_access_groups.length > 0 && ( + + {agent.object_permission.mcp_access_groups.join(", ")} + + )} + {agent.object_permission.mcp_tool_permissions && + Object.keys(agent.object_permission.mcp_tool_permissions).length > 0 && ( + +
+ {Object.entries(agent.object_permission.mcp_tool_permissions).map( + ([serverId, tools]) => ( +
+ {serverId}:{" "} + {Array.isArray(tools) ? tools.join(", ") : String(tools)} +
+ ) + )} +
+ + )} + +
+ )} + {agent.agent_card_params?.skills && agent.agent_card_params.skills.length > 0 && ( diff --git a/ui/litellm-dashboard/src/components/agents/types.ts b/ui/litellm-dashboard/src/components/agents/types.ts index 5c63d334129..2e903b14026 100644 --- a/ui/litellm-dashboard/src/components/agents/types.ts +++ b/ui/litellm-dashboard/src/components/agents/types.ts @@ -1,3 +1,15 @@ +export interface AgentKeyInfo { + key_alias?: string; + token_prefix?: string; + has_key: boolean; +} + +export interface AgentObjectPermission { + mcp_servers?: string[]; + mcp_access_groups?: string[]; + mcp_tool_permissions?: Record; +} + export interface Agent { agent_id: string; agent_name: string; @@ -7,8 +19,10 @@ export interface Agent { }; agent_card_params?: { description?: string; + url?: string; [key: string]: any; }; + object_permission?: AgentObjectPermission; created_at?: string; updated_at?: string; created_by?: string; diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_garden.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_garden.tsx index 703a6ad3e48..d95ae07e7f3 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_garden.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_garden.tsx @@ -1,7 +1,7 @@ import React, { useState } from "react"; import { Input } from "antd"; import { SearchOutlined, ArrowRightOutlined } from "@ant-design/icons"; -import { GuardrailCardInfo, LITELLM_CONTENT_FILTER_CARDS, PARTNER_GUARDRAIL_CARDS, ALL_CARDS } from "./guardrail_garden_data"; +import { GuardrailCardInfo, ALL_CARDS } from "./guardrail_garden_data"; import GuardrailCard from "./guardrail_garden_card"; import GuardrailDetailView from "./guardrail_garden_detail"; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx index 84967f7a83e..404172623b9 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx @@ -194,7 +194,7 @@ const MCPToolConfiguration: React.FC = ({ {filteredTools.length === 0 ? (
- No tools found matching "{toolSearchTerm}" + No tools found matching "{toolSearchTerm}"
) : ( filteredTools.map((tool, index) => ( diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx index 3a572dec893..198bf96c3f0 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx @@ -1,12 +1,12 @@ import React, { useState } from "react"; import { useQuery, useMutation } from "@tanstack/react-query"; import { ToolTestPanel } from "./ToolTestPanel"; -import { MCPTool, MCPToolsViewerProps, MCPContent, CallMCPToolResponse, AUTH_TYPE } from "./types"; +import { MCPTool, MCPToolsViewerProps, MCPContent, CallMCPToolResponse } from "./types"; import { listMCPTools, callMCPTool } from "../networking"; import { Card, Title, Text } from "@tremor/react"; -import { RobotOutlined, ToolOutlined, SearchOutlined, LockOutlined, KeyOutlined } from "@ant-design/icons"; -import { Input, Alert, Button as AntdButton } from "antd"; +import { RobotOutlined, ToolOutlined, SearchOutlined, KeyOutlined } from "@ant-design/icons"; +import { Input, Button as AntdButton } from "antd"; const MCPToolsViewer = ({ serverId, @@ -134,7 +134,7 @@ const MCPToolsViewer = ({ {!showHeaderInput && Object.keys(passthroughHeaders).length === 0 && ( - This server requires additional headers. Click "Configure" to provide values. + This server requires additional headers. Click "Configure" to provide values. )} @@ -255,7 +255,7 @@ const MCPToolsViewer = ({

No tools found

-

No tools match "{toolSearchTerm}"

+

No tools match "{toolSearchTerm}"

) : (
= ({ {!result && !error && complianceResults.length === 0 && (
- Choose a test source above (quick chat or a compliance dataset) and click "Run Test" + Choose a test source above (quick chat or a compliance dataset) and click "Run Test"
)}
diff --git a/ui/litellm-dashboard/src/components/policies/policy_templates.tsx b/ui/litellm-dashboard/src/components/policies/policy_templates.tsx index 6951b368ef3..98ba0acb6a6 100644 --- a/ui/litellm-dashboard/src/components/policies/policy_templates.tsx +++ b/ui/litellm-dashboard/src/components/policies/policy_templates.tsx @@ -1,5 +1,5 @@ import React, { useState, useEffect, useMemo } from "react"; -import { Card, Button, Spin, message, Checkbox, Badge } from "antd"; +import { Card, Button, Spin, message, Checkbox } from "antd"; import { ShieldCheckIcon, ShieldExclamationIcon, diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.test.tsx index 0d746b1f872..c9c5856129f 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.test.tsx @@ -1,4 +1,4 @@ -import { render, screen, waitFor } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; import { KeyInfoHeader, KeyInfoData } from "./KeyInfoHeader"; diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx index 8ff4f53bdd3..f8c588e77a6 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx @@ -4,7 +4,6 @@ import moment from "moment"; import { LogEntry } from "../columns"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import GuardrailViewer from "../GuardrailViewer/GuardrailViewer"; -import CompliancePanel from "../GuardrailViewer/CompliancePanel"; import { CostBreakdownViewer } from "../CostBreakdownViewer"; import { ConfigInfoMessage } from "../ConfigInfoMessage"; import { VectorStoreViewer } from "../VectorStoreViewer"; diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx index 9f199ec8ac9..eb6550d7827 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx @@ -15,7 +15,7 @@ import { fetchAllKeyAliases } from "../key_team_helpers/filter_helpers"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import { PaginatedModelSelect } from "../ModelSelect/PaginatedModelSelect/PaginatedModelSelect"; import FilterComponent, { FilterOption } from "../molecules/filter"; -import { allEndUsersCall, keyInfoV1Call, keyListCall, uiSpendLogsCall } from "../networking"; +import { allEndUsersCall, keyInfoV1Call, uiSpendLogsCall } from "../networking"; import KeyInfoView from "../templates/key_info_view"; import AuditLogs from "./audit_logs"; import { createColumns, LogEntry, type LogsSortField } from "./columns"; diff --git a/ui/litellm-dashboard/src/data/financialCompliancePrompts.ts b/ui/litellm-dashboard/src/data/financialCompliancePrompts.ts index adeacfd9cb5..4c89548e8b8 100644 --- a/ui/litellm-dashboard/src/data/financialCompliancePrompts.ts +++ b/ui/litellm-dashboard/src/data/financialCompliancePrompts.ts @@ -1,7 +1,7 @@ // Auto-generated from block_investment.csv — do not edit manually. // Regenerate: python scripts/generate_compliance_prompts.py --csv ... --output ... -import type { CompliancePrompt, ComplianceFramework } from "./compliancePrompts"; +import type { CompliancePrompt } from "./compliancePrompts"; export const financialCompliancePrompts: CompliancePrompt[] = [ { diff --git a/ui/litellm-dashboard/src/data/insultsCompliancePrompts.ts b/ui/litellm-dashboard/src/data/insultsCompliancePrompts.ts index 4f0948600b8..4162193676c 100644 --- a/ui/litellm-dashboard/src/data/insultsCompliancePrompts.ts +++ b/ui/litellm-dashboard/src/data/insultsCompliancePrompts.ts @@ -1,7 +1,7 @@ // Auto-generated from block_insults.csv — do not edit manually. // Regenerate: python scripts/generate_compliance_prompts.py --csv ... --output ... -import type { CompliancePrompt, ComplianceFramework } from "./compliancePrompts"; +import type { CompliancePrompt } from "./compliancePrompts"; export const insultsCompliancePrompts: CompliancePrompt[] = [ {