feat(agents): assign virtual keys to agents (#22045)

* 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

@@ -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

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,

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?

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')}"
+    )

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<AddAgentFormProps> = ({
   onSuccess,
 }) => {
   const [form] = Form.useForm();
+  const [currentStep, setCurrentStep] = useState(0);
   const [isSubmitting, setIsSubmitting] = useState(false);
   const [agentType, setAgentType] = useState<string>("a2a");
   const [agentTypeMetadata, setAgentTypeMetadata] = useState<AgentCreateInfo[]>([]);
   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<string>("");
+  const [newKeyModels, setNewKeyModels] = useState<string[]>([]);
+  const [existingKeys, setExistingKeys] = useState<any[]>([]);
+  const [selectedExistingKey, setSelectedExistingKey] = useState<string | null>(null);
+  const [loadingKeys, setLoadingKeys] = useState(false);
+
+  // Step 2: results
+  const [createdAgentName, setCreatedAgentName] = useState<string>("");
+  const [createdKeyValue, setCreatedKeyValue] = useState<string | null>(null);
+  const [assignedKeyAlias, setAssignedKeyAlias] = useState<string | null>(null);
+
   // Fetch agent type metadata on mount
   useEffect(() => {
     const fetchMetadata = async () => {
@@ -41,11 +68,87 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
     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<AddAgentFormProps> = ({
 
     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<AddAgentFormProps> = ({
     }
   };
 
-  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<AddAgentFormProps> = ({
     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 = () => (
+    <>
+      <Form.Item
+        label={<span className="text-sm font-medium text-gray-700">Agent Type</span>}
+        required
+        tooltip="Select the type of agent you want to create"
+      >
+        <Select
+          value={agentType}
+          onChange={handleAgentTypeChange}
+          size="large"
+          style={{ width: "100%" }}
+          optionLabelProp="label"
+          dropdownRender={(menu) => (
+            <>
+              {menu}
+              <Divider style={{ margin: "4px 0" }} />
+              <div className="px-2 py-1">
+                <div className="text-xs text-gray-400 font-medium mb-1 uppercase tracking-wide px-2">
+                  Not listed?
+                </div>
+                <div
+                  className={`flex items-center gap-3 px-2 py-2 rounded cursor-pointer transition-colors ${
+                    agentType === CUSTOM_AGENT_TYPE
+                      ? "bg-amber-50"
+                      : "hover:bg-amber-50"
+                  }`}
+                  onClick={() => handleAgentTypeChange(CUSTOM_AGENT_TYPE)}
+                >
+                  <AppstoreOutlined className="text-amber-600 text-lg" />
+                  <div className="flex-1">
+                    <div className="flex items-center gap-2">
+                      <span className="font-medium text-amber-700">Custom / Other</span>
+                      <Tag color="orange" style={{ fontSize: 10, padding: "0 4px" }}>GENERIC</Tag>
+                    </div>
+                    <div className="text-xs text-amber-600">
+                      For agents that don&apos;t follow a standard protocol — just needs a virtual key
+                    </div>
+                  </div>
+                </div>
+              </div>
+            </>
+          )}
+        >
+          {agentTypeMetadata.map((info) => (
+            <Select.Option
+              key={info.agent_type}
+              value={info.agent_type}
+              label={
+                <div className="flex items-center gap-2">
+                  <img src={info.logo_url || ""} alt="" className="w-4 h-4 object-contain" />
+                  <span>{info.agent_type_display_name}</span>
+                </div>
+              }
+            >
+              <div className="flex items-center gap-3 py-1">
+                <img
+                  src={info.logo_url || ""}
+                  alt={info.agent_type_display_name}
+                  className="w-5 h-5 object-contain"
+                />
+                <div>
+                  <div className="font-medium">{info.agent_type_display_name}</div>
+                  {info.description && (
+                    <div className="text-xs text-gray-500">{info.description}</div>
+                  )}
+                </div>
+              </div>
+            </Select.Option>
+          ))}
+        </Select>
+      </Form.Item>
+
+      <div className="mt-4">
+        {agentType === CUSTOM_AGENT_TYPE ? (
+          <div className="space-y-4">
+            <Form.Item
+              label="Agent Name"
+              name="agent_name"
+              rules={[{ required: true, message: "Please enter an agent name" }]}
+            >
+              <Input placeholder="e.g. my-custom-agent" />
+            </Form.Item>
+            <Form.Item
+              label="Description"
+              name="description"
+            >
+              <Input.TextArea placeholder="Describe what this agent does…" rows={3} />
+            </Form.Item>
+          </div>
+        ) : agentType === "a2a" ? (
+          <AgentFormFields showAgentName={true} />
+        ) : selectedAgentTypeInfo?.use_a2a_form_fields ? (
+          <>
+            <AgentFormFields showAgentName={true} />
+            {selectedAgentTypeInfo.credential_fields.length > 0 && (
+              <div className="mt-4 p-4 border border-gray-200 rounded-lg">
+                <h4 className="text-sm font-medium text-gray-700 mb-3">
+                  {selectedAgentTypeInfo.agent_type_display_name} Settings
+                </h4>
+                {selectedAgentTypeInfo.credential_fields.map((field) => (
+                  <Form.Item
+                    key={field.key}
+                    label={field.label}
+                    name={field.key}
+                    rules={
+                      field.required
+                        ? [{ required: true, message: `Please enter ${field.label}` }]
+                        : undefined
+                    }
+                    tooltip={field.tooltip}
+                    initialValue={field.default_value}
+                  >
+                    {field.field_type === "password" ? (
+                      <Input.Password placeholder={field.placeholder || ""} />
+                    ) : (
+                      <Input placeholder={field.placeholder || ""} />
+                    )}
+                  </Form.Item>
+                ))}
+              </div>
+            )}
+          </>
+        ) : selectedAgentTypeInfo ? (
+          <DynamicAgentFormFields agentTypeInfo={selectedAgentTypeInfo} />
+        ) : null}
+      </div>
+    </>
+  );
+
+  const renderAssignKeyStep = () => {
+    const agentName = form.getFieldValue("agent_name") || "your-agent";
+    return (
+      <div>
+        {/* Agent name chip */}
+        <div className="flex justify-center mb-6">
+          <Tag icon={<RobotOutlined />} color="purple" className="px-3 py-1 text-sm">
+            {agentName}
+          </Tag>
+        </div>
+
+        <div className="space-y-3">
+          {/* Option: Create new key */}
+          <div
+            className={`p-4 border-2 rounded-lg cursor-pointer transition-colors ${
+              keyAssignOption === "create_new"
+                ? "border-indigo-600 bg-indigo-50"
+                : "border-gray-200 bg-white hover:border-gray-300"
+            }`}
+            onClick={() => setKeyAssignOption("create_new")}
+          >
+            <div className="flex items-start justify-between">
+              <div className="flex items-start gap-3 flex-1">
+                <Radio
+                  value="create_new"
+                  checked={keyAssignOption === "create_new"}
+                  onChange={() => setKeyAssignOption("create_new")}
+                />
+                <div className="flex-1">
+                  <div className="flex items-center gap-2">
+                    <KeyOutlined className="text-indigo-600" />
+                    <span className="font-medium text-gray-900">Create a new key for this agent</span>
+                  </div>
+                  <p className="text-sm text-gray-500 mt-1">
+                    A dedicated key scoped to this agent.
+                  </p>
+                  {keyAssignOption === "create_new" && (
+                    <div className="mt-3 space-y-3" onClick={(e) => e.stopPropagation()}>
+                      <div>
+                        <label className="text-sm text-gray-600 block mb-1">Key Name</label>
+                        <Input
+                          value={newKeyName}
+                          onChange={(e) => setNewKeyName(e.target.value)}
+                          placeholder="e.g. my-agent-key"
+                        />
+                      </div>
+                      <div>
+                        <label className="text-sm text-gray-600 block mb-1">
+                          Allowed Models <span className="text-gray-400">(optional — leave empty for all models)</span>
+                        </label>
+                        <Select
+                          mode="tags"
+                          style={{ width: "100%" }}
+                          placeholder="e.g. gpt-4o, claude-3-5-sonnet"
+                          value={newKeyModels}
+                          onChange={setNewKeyModels}
+                          tokenSeparators={[","]}
+                        />
+                      </div>
+                    </div>
+                  )}
+                </div>
+              </div>
+              <Tag color="green">Recommended</Tag>
+            </div>
+          </div>
+
+          {/* Option: Assign existing key */}
+          <div
+            className={`p-4 border-2 rounded-lg cursor-pointer transition-colors ${
+              keyAssignOption === "existing_key"
+                ? "border-indigo-600 bg-indigo-50"
+                : "border-gray-200 bg-white hover:border-gray-300"
+            }`}
+            onClick={() => setKeyAssignOption("existing_key")}
+          >
+            <div className="flex items-start gap-3">
+              <Radio
+                value="existing_key"
+                checked={keyAssignOption === "existing_key"}
+                onChange={() => setKeyAssignOption("existing_key")}
+              />
+              <div className="flex-1">
+                <div className="flex items-center gap-2">
+                  <KeyOutlined className="text-gray-500" />
+                  <span className="font-medium text-gray-900">Assign an existing key</span>
+                </div>
+                <p className="text-sm text-gray-500 mt-1">
+                  Re-assign a key you already have to this agent.
+                </p>
+                {keyAssignOption === "existing_key" && (
+                  <div className="mt-3" onClick={(e) => e.stopPropagation()}>
+                    <Select
+                      showSearch
+                      style={{ width: "100%" }}
+                      placeholder="Search by key name…"
+                      loading={loadingKeys}
+                      value={selectedExistingKey}
+                      onChange={(value) => 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,
+                      }))}
+                    />
+                  </div>
+                )}
+              </div>
+            </div>
+          </div>
+        </div>
+
+        <div className="text-center mt-4">
+          <button
+            type="button"
+            className="text-sm text-gray-500 underline hover:text-gray-700"
+            onClick={() => setKeyAssignOption("skip")}
+          >
+            Skip for now — I&apos;ll assign a key later
+          </button>
+        </div>
+      </div>
+    );
+  };
+
+  const renderReadyStep = () => (
+    <div className="text-center py-6">
+      <CheckCircleFilled className="text-5xl text-green-500 mb-4" style={{ fontSize: 48 }} />
+      <h3 className="text-xl font-semibold text-gray-900 mb-2">Agent Created!</h3>
+      <div className="flex justify-center mb-4">
+        <Tag icon={<RobotOutlined />} color="purple" className="px-3 py-1 text-sm">
+          {createdAgentName}
+        </Tag>
+      </div>
+      {createdKeyValue && (
+        <div className="mt-4 text-left max-w-md mx-auto">
+          <CreatedKeyDisplay apiKey={createdKeyValue} />
+        </div>
+      )}
+      {assignedKeyAlias && (
+        <p className="text-sm text-gray-600 mt-2">
+          Key <span className="font-medium">{assignedKeyAlias}</span> has been assigned to this agent.
+        </p>
+      )}
+      {!createdKeyValue && !assignedKeyAlias && keyAssignOption === "skip" && (
+        <p className="text-sm text-gray-500 mt-2">
+          No key assigned. You can create one from the Virtual Keys page.
+        </p>
+      )}
+    </div>
+  );
 
   return (
     <Modal
       title={
         <div className="flex items-center space-x-3 pb-4 border-b border-gray-100">
-          {selectedLogo && (
-            <img
-              src={selectedLogo}
-              alt="Agent"
-              className="w-6 h-6 object-contain"
-            />
+          {selectedLogo && currentStep < 1 && (
+            <img src={selectedLogo} alt="Agent" className="w-6 h-6 object-contain" />
           )}
           <h2 className="text-xl font-semibold text-gray-900">Add New Agent</h2>
         </div>
       }
       open={visible}
-      onCancel={handleCancel}
+      onCancel={handleClose}
       footer={null}
       width={900}
       className="top-8"
@@ -134,103 +534,60 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
       }}
     >
       <div className="mt-4">
+        {/* Step indicator */}
+        <Steps current={currentStep} size="small" className="mb-8">
+          <Step title="Configure" />
+          <Step title="Assign Key" />
+          <Step title="Ready" />
+        </Steps>
+
         <Form
           form={form}
           layout="vertical"
-          onFinish={handleSubmit}
           initialValues={agentType === "a2a" ? getDefaultFormValues() : {}}
           className="space-y-4"
         >
-          {/* Agent Type Selection */}
-          <Form.Item
-            label={<span className="text-sm font-medium text-gray-700">Agent Type</span>}
-            required
-            tooltip="Select the type of agent you want to create"
-          >
-            <Select
-              value={agentType}
-              onChange={handleAgentTypeChange}
-              size="large"
-              style={{ width: "100%" }}
-              optionLabelProp="label"
-            >
-              {agentTypeMetadata.map((info) => (
-                <Select.Option 
-                  key={info.agent_type} 
-                  value={info.agent_type}
-                  label={
-                    <div className="flex items-center gap-2">
-                      <img src={info.logo_url || ""} alt="" className="w-4 h-4 object-contain" />
-                      <span>{info.agent_type_display_name}</span>
-                    </div>
-                  }
-                >
-                  <div className="flex items-center gap-3 py-1">
-                    <img
-                      src={info.logo_url || ""}
-                      alt={info.agent_type_display_name}
-                      className="w-5 h-5 object-contain"
-                    />
-                    <div>
-                      <div className="font-medium">{info.agent_type_display_name}</div>
-                      {info.description && (
-                        <div className="text-xs text-gray-500">{info.description}</div>
-                      )}
-                    </div>
-                  </div>
-                </Select.Option>
-              ))}
-            </Select>
-          </Form.Item>
-
-          {/* Conditional Form Fields */}
-          <div className="mt-6">
-            {agentType === "a2a" ? (
-              <AgentFormFields showAgentName={true} />
-            ) : selectedAgentTypeInfo?.use_a2a_form_fields ? (
-              // A2A-compatible agents (like Pydantic AI) use full A2A form fields
-              // plus any additional credential fields
-              <>
-                <AgentFormFields showAgentName={true} />
-                {selectedAgentTypeInfo.credential_fields.length > 0 && (
-                  <div className="mt-4 p-4 border border-gray-200 rounded-lg">
-                    <h4 className="text-sm font-medium text-gray-700 mb-3">
-                      {selectedAgentTypeInfo.agent_type_display_name} Settings
-                    </h4>
-                    {selectedAgentTypeInfo.credential_fields.map((field) => (
-                      <Form.Item
-                        key={field.key}
-                        label={field.label}
-                        name={field.key}
-                        rules={field.required ? [{ required: true, message: `Please enter ${field.label}` }] : undefined}
-                        tooltip={field.tooltip}
-                        initialValue={field.default_value}
-                      >
-                        {field.field_type === "password" ? (
-                          <Input.Password placeholder={field.placeholder || ""} />
-                        ) : (
-                          <Input placeholder={field.placeholder || ""} />
-                        )}
-                      </Form.Item>
-                    ))}
-                  </div>
-                )}
-              </>
-            ) : selectedAgentTypeInfo ? (
-              <DynamicAgentFormFields agentTypeInfo={selectedAgentTypeInfo} />
-            ) : null}
-          </div>
-
-          {/* Footer Buttons */}
-          <div className="flex items-center justify-end space-x-3 pt-6 border-t border-gray-100 mt-6">
-            <Button variant="secondary" onClick={handleCancel}>
-              Cancel
-            </Button>
-            <Button variant="primary" loading={isSubmitting}>
-              {isSubmitting ? "Creating..." : "Create Agent"}
-            </Button>
-          </div>
+          {currentStep === 0 && renderConfigureStep()}
+          {currentStep === 1 && renderAssignKeyStep()}
+          {currentStep === 2 && renderReadyStep()}
         </Form>
+
+        {/* Footer navigation */}
+        <div className="flex items-center justify-between pt-6 border-t border-gray-100 mt-6">
+          <div>
+            {currentStep > 0 && currentStep < 2 && (
+              <button
+                type="button"
+                onClick={handleBack}
+                className="text-sm text-gray-600 border border-gray-300 rounded px-4 py-2 hover:bg-gray-50"
+              >
+                ← Back
+              </button>
+            )}
+          </div>
+          <div className="flex gap-3">
+            {currentStep < 2 && (
+              <Button variant="secondary" onClick={handleClose}>
+                Cancel
+              </Button>
+            )}
+            {currentStep === 0 && (
+              <Button variant="primary" onClick={handleNext}>
+                Next →
+              </Button>
+            )}
+            {currentStep === 1 && (
+              <Button variant="primary" loading={isSubmitting} onClick={handleCreateAgent}>
+                {isSubmitting ? "Creating..." : "Create Agent →"}
+              </Button>
+            )}
+            {currentStep === 2 && (
+              <Button variant="primary" onClick={handleClose}>
+                Done
+              </Button>
+            )}
+          </div>
+        </div>
       </div>
     </Modal>
   );

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"],

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<AgentFormFieldsProps> = ({ showAgentName = true }) => {
+const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true, visiblePanels }) => {
+  const shouldShow = (key: string) => !visiblePanels || visiblePanels.includes(key);
   return (
     <>
       {showAgentName && (
@@ -32,6 +34,7 @@ const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true
 
       <Collapse defaultActiveKey={['basic']} style={{ marginBottom: 16 }}>
         {/* Basic Information */}
+        {shouldShow(AGENT_FORM_CONFIG.basic.key) && (
         <Panel header={`${AGENT_FORM_CONFIG.basic.title} (Required)`} key={AGENT_FORM_CONFIG.basic.key}>
           {AGENT_FORM_CONFIG.basic.fields.map((field) => (
             <Form.Item
@@ -49,8 +52,10 @@ const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true
             </Form.Item>
           ))}
         </Panel>
+        )}
 
         {/* Skills */}
+        {shouldShow(AGENT_FORM_CONFIG.skills.key) && (
         <Panel header={`${AGENT_FORM_CONFIG.skills.title} (Required)`} key={AGENT_FORM_CONFIG.skills.key}>
           <Form.List name="skills">
             {(fields, { add, remove }) => (
@@ -127,8 +132,10 @@ const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true
             )}
           </Form.List>
         </Panel>
+        )}
 
         {/* Capabilities */}
+        {shouldShow(AGENT_FORM_CONFIG.capabilities.key) && (
         <Panel header={AGENT_FORM_CONFIG.capabilities.title} key={AGENT_FORM_CONFIG.capabilities.key}>
           {AGENT_FORM_CONFIG.capabilities.fields.map((field) => (
             <Form.Item
@@ -141,8 +148,10 @@ const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true
             </Form.Item>
           ))}
         </Panel>
+        )}
 
         {/* Optional Settings */}
+        {shouldShow(AGENT_FORM_CONFIG.optional.key) && (
         <Panel header={AGENT_FORM_CONFIG.optional.title} key={AGENT_FORM_CONFIG.optional.key}>
           {AGENT_FORM_CONFIG.optional.fields.map((field) => (
             <Form.Item
@@ -155,13 +164,17 @@ const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true
             </Form.Item>
           ))}
         </Panel>
+        )}
 
         {/* Cost Configuration */}
+        {shouldShow(AGENT_FORM_CONFIG.cost.key) && (
         <Panel header={AGENT_FORM_CONFIG.cost.title} key={AGENT_FORM_CONFIG.cost.key}>
           <CostConfigFields />
         </Panel>
+        )}
 
         {/* LiteLLM Parameters */}
+        {shouldShow(AGENT_FORM_CONFIG.litellm.key) && (
         <Panel header={AGENT_FORM_CONFIG.litellm.title} key={AGENT_FORM_CONFIG.litellm.key}>
           {AGENT_FORM_CONFIG.litellm.fields.map((field) => (
             <Form.Item
@@ -174,6 +187,7 @@ const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true
             </Form.Item>
           ))}
         </Panel>
+        )}
       </Collapse>
     </>
   );

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,

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<CreateKeyProps> = ({ team, teams, data, addKey }) => {
   const [rotationInterval, setRotationInterval] = useState<string>("30d");
   const [routerSettings, setRouterSettings] = useState<RouterSettingsAccordionValue | null>(null);
   const [routerSettingsKey, setRouterSettingsKey] = useState<number>(0);
+  const [agentsList, setAgentsList] = useState<{ agent_id: string; agent_name: string }[]>([]);
+  const [selectedAgentId, setSelectedAgentId] = useState<string | null>(null);
   const handleOk = () => {
     setIsModalVisible(false);
     form.resetFields();
@@ -180,6 +184,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
     setRotationInterval("30d");
     setRouterSettings(null);
     setRouterSettingsKey((prev) => prev + 1);
+    setSelectedAgentId(null);
   };
 
   const handleCancel = () => {
@@ -195,6 +200,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
     setRotationInterval("30d");
     setRouterSettings(null);
     setRouterSettingsKey((prev) => prev + 1);
+    setSelectedAgentId(null);
   };
 
   useEffect(() => {
@@ -203,6 +209,14 @@ const CreateKey: React.FC<CreateKeyProps> = ({ 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<CreateKeyProps> = ({ 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<CreateKeyProps> = ({ team, teams, data, addKey }) => {
                 <Radio value="you">You</Radio>
                 <Radio value="service_account">Service Account</Radio>
                 {userRole === "Admin" && <Radio value="another_user">Another User</Radio>}
+                <Radio value="agent">Agent <Tag color="purple">New</Tag></Radio>
               </Radio.Group>
             </Form.Item>
 
@@ -583,6 +604,32 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
                 </div>
               </Form.Item>
             )}
+            {keyOwner === "agent" && (
+              <div className="mt-4 p-4 bg-purple-50 border border-purple-200 rounded-md">
+                <div className="mb-3">
+                  <span className="text-sm font-medium text-gray-700">
+                    Select Agent <span className="text-red-500">*</span>
+                  </span>
+                </div>
+                <Select
+                  showSearch
+                  placeholder="Select an agent"
+                  style={{ width: "100%" }}
+                  value={selectedAgentId}
+                  onChange={(value) => setSelectedAgentId(value)}
+                  filterOption={(input, option) =>
+                    (option?.label as string)?.toLowerCase().includes(input.toLowerCase())
+                  }
+                  options={agentsList.map((a) => ({
+                    label: a.agent_name || a.agent_id,
+                    value: a.agent_id,
+                  }))}
+                />
+                <div className="text-xs text-gray-500 mt-2">
+                  This key will be used by the selected agent to make requests to LiteLLM
+                </div>
+              </div>
+            )}
             <Form.Item
               label={
                 <span>
@@ -1380,35 +1427,9 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
         <Modal open={isModalVisible} onOk={handleOk} onCancel={handleCancel} footer={null}>
           <Grid numItems={1} className="gap-2 w-full">
             <Title>Save your Key</Title>
-            <Col numColSpan={1}>
-              <p>
-                Please save this secret key somewhere safe and accessible. For security reasons,{" "}
-                <b>you will not be able to view it again</b> through your LiteLLM account. If you lose this secret key,
-                you will need to generate a new one.
-              </p>
-            </Col>
             <Col numColSpan={1}>
               {apiKey != null ? (
-                <div>
-                  <Text className="mt-3">Virtual Key:</Text>
-                  <div
-                    style={{
-                      background: "#f8f8f8",
-                      padding: "10px",
-                      borderRadius: "5px",
-                      marginBottom: "10px",
-                    }}
-                  >
-                    <pre style={{ wordWrap: "break-word", whiteSpace: "normal" }}>{apiKey}</pre>
-                  </div>
-
-                  <CopyToClipboard text={apiKey} onCopy={handleCopy}>
-                    <Button className="mt-3">Copy Virtual Key</Button>
-                  </CopyToClipboard>
-                  {/* <Button className="mt-3" onClick={sendSlackAlert}>
-                    Test Key
-                </Button> */}
-                </div>
+                <CreatedKeyDisplay apiKey={apiKey} />
               ) : (
                 <Text>Key being created, this might take 30s</Text>
               )}

ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx

@@ -0,0 +1,53 @@
+import React, { useState } from "react";
+import { CopyToClipboard } from "react-copy-to-clipboard";
+import { Button, message } from "antd";
+
+interface CreatedKeyDisplayProps {
+  apiKey: string;
+}
+
+/**
+ * Shared component for displaying a newly-created virtual key.
+ * Used on the Virtual Keys page and in the Add Agent wizard.
+ */
+const CreatedKeyDisplay: React.FC<CreatedKeyDisplayProps> = ({ apiKey }) => {
+  const [copied, setCopied] = useState(false);
+
+  const handleCopy = () => {
+    setCopied(true);
+    message.success("Key copied to clipboard");
+    setTimeout(() => setCopied(false), 2000);
+  };
+
+  return (
+    <div>
+      <p className="mb-2">
+        Please save this secret key somewhere safe and accessible. For security reasons,{" "}
+        <b>you will not be able to view it again</b> through your LiteLLM account. If you
+        lose this secret key, you will need to generate a new one.
+      </p>
+
+      <p className="text-sm text-gray-600 mt-3 mb-1">Virtual Key:</p>
+      <div
+        style={{
+          background: "#f8f8f8",
+          padding: "10px",
+          borderRadius: "5px",
+          marginBottom: "10px",
+        }}
+      >
+        <pre style={{ wordWrap: "break-word", whiteSpace: "normal", margin: 0 }}>
+          {apiKey}
+        </pre>
+      </div>
+
+      <CopyToClipboard text={apiKey} onCopy={handleCopy}>
+        <Button type="primary" style={{ marginTop: 12 }}>
+          {copied ? "Copied!" : "Copy Virtual Key"}
+        </Button>
+      </CopyToClipboard>
+    </div>
+  );
+};
+
+export default CreatedKeyDisplay;