From 60bcb26dc8a8c064183ca516aff54172bb0b0522 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 24 Feb 2026 18:28:16 -0800 Subject: [PATCH] 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 75b9f91acd..4053d9d077 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 a230c2e933..c1165ab26d 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 23917cf7c7..8f746b1f9c 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 ffb4e95542..05df3c2dcb 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 f4e0137bd0..360b9b7ee9 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 9dd41eed8a..f85c4daac6 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 4dc4ad6829..d5429d2a3b 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 a3917c2f09..3aebcb2733 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 9a99870c26..714507da1a 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 * + +
+