From 57a48e352695d5d3343813907482711aec6f6f5d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 10 Mar 2026 21:03:20 -0700 Subject: [PATCH 1/2] fix(agents.tsx): support granting agents access to subagents --- .../litellm_proxy_extras/schema.prisma | 1 + litellm/proxy/_types.py | 75 ++-- litellm/proxy/schema.prisma | 1 + litellm/types/agents.py | 2 + schema.prisma | 1 + .../src/components/agents/add_agent_form.tsx | 373 ++++++++++-------- 6 files changed, 252 insertions(+), 201 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 8d4bdffb2d..939f1eb0f4 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -267,6 +267,7 @@ model LiteLLM_ObjectPermissionTable { vector_stores String[] @default([]) agents String[] @default([]) agent_access_groups String[] @default([]) + models String[] @default([]) blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 36790e9fea..add3ab4a1f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,59 +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 ( - MCPAuthType, - MCPCredentials, - MCPTransport, - MCPTransportType, -) +from litellm.types.llms.openai import (AllMessageValues, OpenAIFileObject, + ResponsesAPIResponse) +from litellm.types.mcp import (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 @@ -855,6 +836,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): vector_stores: Optional[List[str]] = None agents: Optional[List[str]] = None agent_access_groups: Optional[List[str]] = None + models: Optional[List[str]] = None class GenerateRequestBase(LiteLLMPydanticObjectBase): @@ -2470,7 +2452,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, @@ -2502,7 +2485,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, @@ -2908,7 +2892,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/schema.prisma b/litellm/proxy/schema.prisma index 721c3e404d..b68872e2ed 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -267,6 +267,7 @@ model LiteLLM_ObjectPermissionTable { vector_stores String[] @default([]) agents String[] @default([]) agent_access_groups String[] @default([]) + models String[] @default([]) blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 951fbfcabd..efb2e73bfb 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -172,6 +172,8 @@ class AgentObjectPermission(TypedDict, total=False): mcp_servers: Optional[List[str]] mcp_access_groups: Optional[List[str]] mcp_tool_permissions: Optional[Dict[str, List[str]]] + models: Optional[List[str]] + agents: Optional[List[str]] class AgentConfig(TypedDict, total=False): diff --git a/schema.prisma b/schema.prisma index 8d4bdffb2d..939f1eb0f4 100644 --- a/schema.prisma +++ b/schema.prisma @@ -267,6 +267,7 @@ model LiteLLM_ObjectPermissionTable { vector_stores String[] @default([]) agents String[] @default([]) agent_access_groups String[] @default([]) + models String[] @default([]) blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] 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 0cec0331f4..5b739e5a16 100644 --- a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx @@ -6,6 +6,7 @@ import CreatedKeyDisplay from "../shared/CreatedKeyDisplay"; import { createAgentCall, getAgentCreateMetadata, + getAgentsList, keyCreateForAgentCall, keyListCall, keyUpdateCall, @@ -45,7 +46,7 @@ const AddAgentForm: React.FC = ({ const [agentTypeMetadata, setAgentTypeMetadata] = useState([]); const [loadingMetadata, setLoadingMetadata] = useState(false); - // Step 1: key assignment state + // Step 3: key assignment state const [keyAssignOption, setKeyAssignOption] = useState<"create_new" | "existing_key" | "skip">("create_new"); const [newKeyName, setNewKeyName] = useState(""); const [newKeyModels, setNewKeyModels] = useState([]); @@ -54,8 +55,10 @@ const AddAgentForm: React.FC = ({ const [loadingKeys, setLoadingKeys] = useState(false); const [availableModels, setAvailableModels] = useState([]); const [loadingModels, setLoadingModels] = useState(false); + const [availableAgents, setAvailableAgents] = useState<{agent_id: string; agent_name: string}[]>([]); + const [loadingAgents, setLoadingAgents] = useState(false); - // Step 2: results + // Step 4: results const [createdAgentName, setCreatedAgentName] = useState(""); const [createdKeyValue, setCreatedKeyValue] = useState(null); const [assignedKeyAlias, setAssignedKeyAlias] = useState(null); @@ -82,9 +85,9 @@ const AddAgentForm: React.FC = ({ fetchMetadata(); }, []); - // Fetch existing keys when assign key step becomes active (step 2) + // Fetch existing keys when Agent Management step becomes active (step 3) useEffect(() => { - if (currentStep === 2 && accessToken && existingKeys.length === 0) { + if (currentStep === 3 && accessToken && existingKeys.length === 0) { const fetchKeys = async () => { setLoadingKeys(true); try { @@ -100,9 +103,9 @@ const AddAgentForm: React.FC = ({ } }, [currentStep, accessToken]); - // Fetch available models when Assign Key step is active (same list as key generation) + // Fetch available models when Agent Management step is active (same list as key generation) useEffect(() => { - if (currentStep !== 2 || !accessToken || !userId || !userRole) return; + if ((currentStep !== 1 && currentStep !== 3) || !accessToken || !userId || !userRole) return; let cancelled = false; setLoadingModels(true); modelAvailableCall(accessToken, userId, userRole) @@ -125,6 +128,25 @@ const AddAgentForm: React.FC = ({ }; }, [currentStep, accessToken, userId, userRole]); + useEffect(() => { + if (currentStep !== 1 || !accessToken) return; + let cancelled = false; + setLoadingAgents(true); + getAgentsList(accessToken) + .then((response) => { + if (cancelled) return; + const agents = response?.agents ?? []; + setAvailableAgents(agents.map((a: any) => ({ agent_id: a.agent_id, agent_name: a.agent_name }))); + }) + .catch((error) => { + if (!cancelled) console.error("Error fetching agents:", error); + }) + .finally(() => { + if (!cancelled) setLoadingAgents(false); + }); + return () => { cancelled = true; }; + }, [currentStep, accessToken]); + const selectedAgentTypeInfo = agentTypeMetadata.find( (info) => info.agent_type === agentType ); @@ -207,11 +229,14 @@ const AddAgentForm: React.FC = ({ // Build object_permission from MCP Tools step (allowed_mcp_servers_and_groups, mcp_tool_permissions) const mcpServersAndGroups = values.allowed_mcp_servers_and_groups; const mcpToolPermissions = values.mcp_tool_permissions || {}; - if ( - mcpServersAndGroups && - (mcpServersAndGroups.servers?.length > 0 || mcpServersAndGroups.accessGroups?.length > 0) || - Object.keys(mcpToolPermissions).length > 0 - ) { + const entitlementModels = values.entitlement_models || []; + const entitlementAgents = values.entitlement_agents || []; + const hasObjectPermission = + (mcpServersAndGroups?.servers?.length > 0 || mcpServersAndGroups?.accessGroups?.length > 0) || + Object.keys(mcpToolPermissions).length > 0 || + entitlementModels.length > 0 || + entitlementAgents.length > 0; + if (hasObjectPermission) { agentData.object_permission = {}; if (mcpServersAndGroups?.servers?.length > 0) { agentData.object_permission.mcp_servers = mcpServersAndGroups.servers; @@ -222,6 +247,12 @@ const AddAgentForm: React.FC = ({ if (Object.keys(mcpToolPermissions).length > 0) { agentData.object_permission.mcp_tool_permissions = mcpToolPermissions; } + if (entitlementModels.length > 0) { + agentData.object_permission.models = entitlementModels; + } + if (entitlementAgents.length > 0) { + agentData.object_permission.agents = entitlementAgents; + } } // Wire trace-id flags and budget controls into agent litellm_params (before create call) @@ -264,7 +295,7 @@ const AddAgentForm: React.FC = ({ setAssignedKeyAlias(keyInfo?.key_alias || selectedExistingKey.slice(0, 12) + "…"); } - setCurrentStep(3); + setCurrentStep(4); onSuccess(); } catch (error) { console.error("Error creating agent:", error); @@ -293,11 +324,54 @@ const AddAgentForm: React.FC = ({ onClose(); }; - const renderMCPToolsStep = () => ( + const renderEntitlementsStep = () => (

- Optionally restrict which MCP servers and tools this agent can use. Leave empty to allow all (subject to key/team permissions). + Configure which models, agents, and MCP tools this agent is allowed to use. Leave fields empty to allow all (subject to key/team permissions).

+ + Allowed Models} + name="entitlement_models" + tooltip="Restrict which models this agent can call. Leave empty to allow all." + > + + (option?.label as string ?? "").toLowerCase().includes(input.toLowerCase()) + } + options={availableAgents.map((a) => ({ + label: a.agent_name, + value: a.agent_id, + }))} + /> + + + + @@ -338,122 +412,121 @@ const AddAgentForm: React.FC = ({
)} + + ); - Tracing, - children: ( -
-
-
- - Require x-litellm-trace-id on calls TO this agent - -

- Only accept this agent being invoked with a trace-id (e.g. when used as a sub-agent). -

-
- -
- -
-
- - Require x-litellm-trace-id on calls BY this agent - -

- Requires LLM/MCP calls made by this agent to include x-litellm-trace-id for session tracking. -

-
- { - setRequireTraceIdOutbound(checked); - if (!checked) { - setMaxIterations(null); - setMaxBudgetPerSession(null); - } - }} - /> -
-
- ), - }, - { - key: "budgets_and_rate_limits", - label: Budgets & Rate Limits, - children: ( -
- {!requireTraceIdOutbound && ( -
- Enable "Require x-litellm-trace-id on calls BY this agent" in Tracing to configure budgets and rate limits. -
- )} - -
Session Budgets
-
-
- - setMaxIterations(val)} - /> -

Hard cap on LLM calls per session

-
-
- - setMaxBudgetPerSession(val)} - /> -

Max spend per trace before returning 429

-
-
- - - -
Agent Rate Limits
-

- Global rate limits applied across all callers of this agent. + const renderObservabilityStep = () => ( +

+
+

Tracing

+
+
+
+ + Require x-litellm-trace-id on calls TO this agent + +

+ Only accept this agent being invoked with a trace-id (e.g. when used as a sub-agent).

-
- - - - - - -
- -
Per-Session Rate Limits
-

- Rate limits per session (x-litellm-trace-id). Each session gets its own counters. -

-
- - - - - - -
- ), - }, - ]} /> + +
+ +
+
+ + Require x-litellm-trace-id on calls BY this agent + +

+ Requires LLM/MCP calls made by this agent to include x-litellm-trace-id for session tracking. +

+
+ { + setRequireTraceIdOutbound(checked); + if (!checked) { + setMaxIterations(null); + setMaxBudgetPerSession(null); + } + }} + /> +
+
+
+ + + +
+

Budgets & Rate Limits

+
+ {!requireTraceIdOutbound && ( +
+ Enable "Require x-litellm-trace-id on calls BY this agent" in Tracing to configure budgets and rate limits. +
+ )} + +
Session Budgets
+
+
+ + setMaxIterations(val)} + /> +

Hard cap on LLM calls per session

+
+
+ + setMaxBudgetPerSession(val)} + /> +

Max spend per trace before returning 429

+
+
+ + + +
Agent Rate Limits
+

+ Global rate limits applied across all callers of this agent. +

+
+ + + + + + +
+ +
Per-Session Rate Limits
+

+ Rate limits per session (x-litellm-trace-id). Each session gets its own counters. +

+
+ + + + + + +
+
+
); @@ -645,25 +718,6 @@ const AddAgentForm: React.FC = ({ placeholder="e.g. my-agent-key" />
-
- -