Merge pull request #23595 from BerriAI/litellm_ui_keys_org_13

[Feature] UI - Keys: Add Organization Dropdown to Create/Edit Key

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -2301,23 +2301,16 @@ async def validate_key_team_change(
     # Check if the team has access to the key's models
     if len(key.models) > 0:
         for model in key.models:
+            # Skip special sentinel values — "all-team-models" means
+            # "use whatever the team allows", so it's always valid.
+            if model == SpecialModelNames.all_team_models.value:
+                continue
             await can_team_access_model(
                 model=model,
                 team_object=team,
                 llm_router=llm_router,
             )
 
-    # Check if the key's user_id is a member of the team
-    member_object = _get_user_in_team(
-        team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id
-    )
-    if key.user_id is not None:
-        if not member_object:
-            raise HTTPException(
-                status_code=403,
-                detail=f"User={key.user_id} is not a member of the team={team.team_id}. Check team members via `/team/info`.",
-            )
-
     # Check if the key's tpm/rpm limit is less than the team's tpm/rpm limit
     if key.tpm_limit is not None:
         if team.tpm_limit and key.tpm_limit > team.tpm_limit:
@@ -2331,6 +2324,17 @@ async def validate_key_team_change(
                 detail=f"Key={key.token} has a rpm_limit={key.rpm_limit} which is greater than the team's rpm_limit={team.rpm_limit}.",
             )
 
+    # Check if the key's user_id is a member of the team
+    member_object = _get_user_in_team(
+        team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id
+    )
+    if key.user_id is not None:
+        if not member_object:
+            raise HTTPException(
+                status_code=403,
+                detail=f"User={key.user_id} is not a member of the team={team.team_id}. Check team members via `/team/info`.",
+            )
+
     # Check if the person initiating the change is a Proxy Admin or Team Admin
     if change_initiated_by.user_role == LitellmUserRoles.PROXY_ADMIN.value:
         return

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

@@ -1502,6 +1502,57 @@ async def test_validate_key_team_change_with_member_permissions():
                     )
 
 
+@pytest.mark.asyncio
+async def test_validate_key_team_change_skips_all_team_models_sentinel():
+    """
+    Test that validate_key_team_change skips the 'all-team-models' sentinel
+    value when checking if the target team can access the key's models.
+
+    Keys with models=["all-team-models"] mean "use whatever models the team
+    allows", so moving them to any team should not fail model validation.
+    """
+    from unittest.mock import AsyncMock, MagicMock, patch
+
+    mock_key = MagicMock()
+    mock_key.user_id = "test-user-123"
+    mock_key.models = ["all-team-models"]
+    mock_key.tpm_limit = None
+    mock_key.rpm_limit = None
+
+    mock_team = MagicMock()
+    mock_team.team_id = "test-team-456"
+    mock_team.models = ["gpt-4", "claude-3"]
+    mock_team.members_with_roles = []
+    mock_team.tpm_limit = None
+    mock_team.rpm_limit = None
+
+    mock_change_initiator = MagicMock()
+    mock_change_initiator.user_id = "test-user-123"
+    mock_change_initiator.user_role = LitellmUserRoles.PROXY_ADMIN.value
+
+    mock_router = MagicMock()
+
+    with patch(
+        "litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model",
+        new_callable=AsyncMock,
+    ) as mock_can_access:
+        with patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team"
+        ) as mock_get_user:
+            mock_get_user.return_value = MagicMock()
+
+            await validate_key_team_change(
+                key=mock_key,
+                team=mock_team,
+                change_initiated_by=mock_change_initiator,
+                llm_router=mock_router,
+            )
+
+            # can_team_access_model should NOT have been called since
+            # "all-team-models" is a sentinel that should be skipped
+            mock_can_access.assert_not_called()
+
+
 def test_key_rotation_fields_helper():
     """
     Test the key data update logic for rotation fields.

ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx

@@ -0,0 +1,69 @@
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { describe, expect, it, vi } from "vitest";
+import OrganizationDropdown from "./OrganizationDropdown";
+
+const MOCK_ORGS = [
+  {
+    organization_id: "org-1",
+    organization_alias: "Engineering",
+    budget_id: "",
+    metadata: {},
+    models: [],
+    spend: 0,
+    model_spend: {},
+    created_at: "",
+    created_by: "",
+    updated_at: "",
+  },
+  {
+    organization_id: "org-2",
+    organization_alias: "Sales",
+    budget_id: "",
+    metadata: {},
+    models: [],
+    spend: 0,
+    model_spend: {},
+    created_at: "",
+    created_by: "",
+    updated_at: "",
+  },
+];
+
+describe("OrganizationDropdown", () => {
+  it("should render", () => {
+    render(<OrganizationDropdown organizations={MOCK_ORGS} />);
+    expect(screen.getByRole("combobox")).toBeInTheDocument();
+  });
+
+  it("should display organization options when opened", async () => {
+    const user = userEvent.setup();
+    render(<OrganizationDropdown organizations={MOCK_ORGS} />);
+
+    await user.click(screen.getByRole("combobox"));
+
+    expect(await screen.findByText("Engineering")).toBeInTheDocument();
+    expect(screen.getByText("Sales")).toBeInTheDocument();
+  });
+
+  it("should call onChange with the org id when an organization is selected", async () => {
+    const onChange = vi.fn();
+    const user = userEvent.setup();
+    render(<OrganizationDropdown organizations={MOCK_ORGS} onChange={onChange} />);
+
+    await user.click(screen.getByRole("combobox"));
+    await user.click(await screen.findByText("Engineering"));
+
+    expect(onChange).toHaveBeenCalledWith("org-1", expect.anything());
+  });
+
+  it("should add ant-select-disabled class when disabled prop is true", () => {
+    const { container } = render(<OrganizationDropdown organizations={MOCK_ORGS} disabled={true} />);
+    expect(container.querySelector(".ant-select-disabled")).toBeTruthy();
+  });
+
+  it("should render with empty organizations list", () => {
+    render(<OrganizationDropdown organizations={[]} />);
+    expect(screen.getByRole("combobox")).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx

@@ -0,0 +1,52 @@
+import React from "react";
+import { Select } from "antd";
+import { Organization } from "../networking";
+
+interface OrganizationDropdownProps {
+  organizations?: Organization[] | null;
+  value?: string;
+  onChange?: (value: string) => void;
+  disabled?: boolean;
+  loading?: boolean;
+}
+
+const OrganizationDropdown: React.FC<OrganizationDropdownProps> = ({
+  organizations,
+  value,
+  onChange,
+  disabled,
+  loading,
+}) => {
+  return (
+    <Select
+      showSearch
+      placeholder="Search or select an organization"
+      value={value}
+      onChange={onChange}
+      disabled={disabled}
+      loading={loading}
+      allowClear
+      filterOption={(input, option) => {
+        if (!option) return false;
+        const org = organizations?.find((o) => o.organization_id === option.key);
+        if (!org) return false;
+
+        const searchTerm = input.toLowerCase().trim();
+        const orgAlias = (org.organization_alias || "").toLowerCase();
+        const orgId = (org.organization_id || "").toLowerCase();
+
+        return orgAlias.includes(searchTerm) || orgId.includes(searchTerm);
+      }}
+      optionFilterProp="children"
+    >
+      {organizations?.map((org) => (
+        <Select.Option key={org.organization_id} value={org.organization_id}>
+          <span className="font-medium">{org.organization_alias}</span>{" "}
+          <span className="text-gray-500">({org.organization_id})</span>
+        </Select.Option>
+      ))}
+    </Select>
+  );
+};
+
+export default OrganizationDropdown;

ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx

@@ -213,7 +213,22 @@ vi.mock("../common_components/PassThroughRoutesSelector", () => ({ default: () =
 vi.mock("../common_components/PremiumLoggingSettings", () => ({ default: () => null }));
 vi.mock("../common_components/RateLimitTypeFormItem", () => ({ default: () => null }));
 vi.mock("../common_components/RouterSettingsAccordion", () => ({ default: () => null }));
-vi.mock("../common_components/team_dropdown", () => ({ default: () => null }));
+vi.mock("../common_components/team_dropdown", () => ({
+  default: ({ teams, onChange, disabled }: { teams?: any[]; onChange?: (v: string) => void; disabled?: boolean }) => (
+    <select
+      data-testid="team-dropdown"
+      disabled={disabled}
+      onChange={(e) => onChange?.(e.target.value)}
+    >
+      <option value="">Select team</option>
+      {teams?.map((t: any) => (
+        <option key={t.team_id} value={t.team_id}>
+          {t.team_alias}
+        </option>
+      ))}
+    </select>
+  ),
+}));
 vi.mock("../CreateUserButton", () => ({ CreateUserButton: () => null }));
 vi.mock("../mcp_server_management/MCPServerSelector", () => ({ default: () => null }));
 vi.mock("../mcp_server_management/MCPToolPermissions", () => ({ default: () => null }));
@@ -227,6 +242,31 @@ vi.mock("@/app/(dashboard)/hooks/projects/useProjects", () => ({
   useProjects: vi.fn().mockReturnValue({ data: [], isLoading: false }),
 }));
 
+vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({
+  useOrganizations: vi.fn().mockReturnValue({
+    data: [
+      { organization_id: "org-1", organization_alias: "Engineering" },
+      { organization_id: "org-2", organization_alias: "Sales" },
+    ],
+    isLoading: false,
+  }),
+}));
+
+vi.mock("../common_components/OrganizationDropdown", () => ({
+  default: ({ value, onChange, disabled }: { value?: string; onChange?: (v: string) => void; disabled?: boolean }) => (
+    <select
+      data-testid="org-dropdown"
+      disabled={disabled}
+      value={value || ""}
+      onChange={(e) => onChange?.(e.target.value)}
+    >
+      <option value="">Select org</option>
+      <option value="org-1">Engineering</option>
+      <option value="org-2">Sales</option>
+    </select>
+  ),
+}));
+
 vi.mock("../common_components/ProjectDropdown", () => ({
   default: ({ value, onChange }: { value?: string; onChange?: (v: string) => void }) => (
     <input
@@ -408,4 +448,81 @@ describe("CreateKey", () => {
       expect(setFieldsValueMock).toHaveBeenCalledWith({ key_type: "management" });
     });
   });
+
+  describe("organization dropdown", () => {
+    it("should render the organization dropdown when modal is open", async () => {
+      renderWithProviders(<CreateKey {...defaultProps} />);
+
+      act(() => {
+        fireEvent.click(screen.getByRole("button", { name: /create new key/i }));
+      });
+
+      await waitFor(() => {
+        expect(screen.getByTestId("org-dropdown")).toBeInTheDocument();
+      });
+    });
+
+    it("should disable the organization dropdown for non-admin users", async () => {
+      authorizedState = { ...defaultAuthorizedState, userRole: "Internal User" };
+
+      renderWithProviders(<CreateKey {...defaultProps} />);
+
+      act(() => {
+        fireEvent.click(screen.getByRole("button", { name: /create new key/i }));
+      });
+
+      await waitFor(() => {
+        expect(screen.getByTestId("org-dropdown")).toBeDisabled();
+      });
+    });
+
+    it("should enable the organization dropdown for admin users", async () => {
+      authorizedState = { ...defaultAuthorizedState, userRole: "Admin" };
+
+      renderWithProviders(<CreateKey {...defaultProps} />);
+
+      act(() => {
+        fireEvent.click(screen.getByRole("button", { name: /create new key/i }));
+      });
+
+      await waitFor(() => {
+        expect(screen.getByTestId("org-dropdown")).not.toBeDisabled();
+      });
+    });
+
+    it("should render team dropdown alongside organization dropdown", async () => {
+      const teamsWithOrg = [
+        { team_id: "team-1", team_alias: "Team Alpha", organization_id: "org-1", models: [] },
+      ];
+
+      renderWithProviders(<CreateKey {...defaultProps} teams={teamsWithOrg as any} />);
+
+      act(() => {
+        fireEvent.click(screen.getByRole("button", { name: /create new key/i }));
+      });
+
+      await waitFor(() => {
+        expect(screen.getByTestId("org-dropdown")).toBeInTheDocument();
+        expect(screen.getByTestId("team-dropdown")).toBeInTheDocument();
+      });
+    });
+
+    it("should set organization_id in form state when org is selected", async () => {
+      renderWithProviders(<CreateKey {...defaultProps} />);
+
+      act(() => {
+        fireEvent.click(screen.getByRole("button", { name: /create new key/i }));
+      });
+
+      await waitFor(() => {
+        expect(screen.getByTestId("org-dropdown")).toBeInTheDocument();
+      });
+
+      act(() => {
+        fireEvent.change(screen.getByTestId("org-dropdown"), { target: { value: "org-1" } });
+      });
+
+      expect(formStateRef.current["organization_id"]).toBe("org-1");
+    });
+  });
 });

ui/litellm-dashboard/src/components/organisms/create_key_button.tsx

@@ -1,5 +1,6 @@
 "use client";
 import { keyKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
+import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations";
 import { useProjects } from "@/app/(dashboard)/hooks/projects/useProjects";
 import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
@@ -23,6 +24,7 @@ import PremiumLoggingSettings from "../common_components/PremiumLoggingSettings"
 import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem";
 import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "../common_components/RouterSettingsAccordion";
 import TeamDropdown from "../common_components/team_dropdown";
+import OrganizationDropdown from "../common_components/OrganizationDropdown";
 import ProjectDropdown from "../common_components/ProjectDropdown";
 import { CreateUserButton } from "../CreateUserButton";
 import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
@@ -160,6 +162,7 @@ export const fetchUserModels = async (
 const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOpenCreate, prefillData }) => {
   const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized();
   const canEditGuardrails = premiumUser || (userRole != null && rolesWithWriteAccess.includes(userRole));
+  const { data: organizations, isLoading: isOrganizationsLoading } = useOrganizations();
   const { data: projects, isLoading: isProjectsLoading } = useProjects();
   const { data: uiSettingsData } = useUISettings();
   const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui);
@@ -179,6 +182,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
   const [promptsList, setPromptsList] = useState<string[]>([]);
   const [loggingSettings, setLoggingSettings] = useState<any[]>([]);
   const [selectedCreateKeyTeam, setSelectedCreateKeyTeam] = useState<Team | null>(team);
+  const [selectedOrganizationId, setSelectedOrganizationId] = useState<string | null>(null);
   const [selectedProjectId, setSelectedProjectId] = useState<string | null>(null);
   const [isCreateUserModalVisible, setIsCreateUserModalVisible] = useState(false);
   const [newlyCreatedUserId, setNewlyCreatedUserId] = useState<string | null>(null);
@@ -207,6 +211,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
     setRouterSettings(null);
     setRouterSettingsKey((prev) => prev + 1);
     setSelectedAgentId(null);
+    setSelectedOrganizationId(null);
     setSelectedProjectId(null);
   };
 
@@ -224,6 +229,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
     setRouterSettings(null);
     setRouterSettingsKey((prev) => prev + 1);
     setSelectedAgentId(null);
+    setSelectedOrganizationId(null);
     setSelectedProjectId(null);
   };
 
@@ -752,6 +758,32 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
                 </div>
               </div>
             )}
+            <Form.Item
+              label={
+                <span>
+                  Organization{" "}
+                  <Tooltip title="The organization this key belongs to. Selecting an organization filters the available teams.">
+                    <InfoCircleOutlined style={{ marginLeft: "4px" }} />
+                  </Tooltip>
+                </span>
+              }
+              name="organization_id"
+              className="mt-4"
+            >
+              <OrganizationDropdown
+                organizations={organizations}
+                loading={isOrganizationsLoading}
+                disabled={userRole !== "Admin"}
+                onChange={(orgId) => {
+                  setSelectedOrganizationId(orgId || null);
+                  // Clear team and project when org changes
+                  setSelectedCreateKeyTeam(null);
+                  setSelectedProjectId(null);
+                  form.setFieldValue("team_id", undefined);
+                  form.setFieldValue("project_id", undefined);
+                }}
+              />
+            </Form.Item>
             <Form.Item
               label={
                 <span>
@@ -773,7 +805,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
               help={keyOwner === "service_account" ? "required" : ""}
             >
               <TeamDropdown
-                teams={teams}
+                teams={selectedOrganizationId ? teams?.filter((t) => t.organization_id === selectedOrganizationId) : teams}
                 disabled={selectedProjectId !== null}
                 loading={!teams}
                 onChange={(teamId) => {
@@ -781,6 +813,14 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
                   setSelectedCreateKeyTeam(selectedTeam);
                   setSelectedProjectId(null);
                   form.setFieldValue("project_id", undefined);
+                  // Auto-populate org from team for non-admin users
+                  if (selectedTeam?.organization_id) {
+                    setSelectedOrganizationId(selectedTeam.organization_id);
+                    form.setFieldValue("organization_id", selectedTeam.organization_id);
+                  } else if (!teamId) {
+                    setSelectedOrganizationId(null);
+                    form.setFieldValue("organization_id", undefined);
+                  }
                 }}
               />
             </Form.Item>
@@ -1531,6 +1571,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
                         excludedFields={[
                           "key_alias",
                           "team_id",
+                          "organization_id",
                           "models",
                           "duration",
                           "metadata",

ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx

@@ -53,6 +53,16 @@ vi.mock("../organisms/create_key_button", () => ({
   fetchTeamModels: vi.fn().mockResolvedValue(["team-model-1", "team-model-2"]),
 }));
 
+vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({
+  useOrganizations: vi.fn().mockReturnValue({
+    data: [
+      { organization_id: "org-1", organization_alias: "Engineering" },
+      { organization_id: "org-2", organization_alias: "Sales" },
+    ],
+    isLoading: false,
+  }),
+}));
+
 vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({
   useAccessGroups: vi.fn().mockReturnValue({
     data: [
@@ -576,4 +586,91 @@ describe("KeyEditView", () => {
       resolveSubmit();
     }
   });
+
+  describe("organization dropdown", () => {
+    it("should render the organization dropdown", async () => {
+      renderWithProviders(
+        <KeyEditView
+          keyData={MOCK_KEY_DATA}
+          onCancel={() => {}}
+          onSubmit={async () => {}}
+          accessToken=""
+          userID=""
+          userRole="Admin"
+          premiumUser={false}
+        />,
+      );
+
+      await waitFor(() => {
+        expect(screen.getByText("Organization")).toBeInTheDocument();
+      });
+    });
+
+    it("should disable the organization dropdown for non-admin users", async () => {
+      const { container } = renderWithProviders(
+        <KeyEditView
+          keyData={MOCK_KEY_DATA}
+          onCancel={() => {}}
+          onSubmit={async () => {}}
+          accessToken=""
+          userID=""
+          userRole="Internal User"
+          premiumUser={false}
+        />,
+      );
+
+      await waitFor(() => {
+        expect(screen.getByText("Organization")).toBeInTheDocument();
+      });
+
+      const orgFormItem = screen.getByText("Organization").closest(".ant-form-item");
+      const disabledSelect = orgFormItem?.querySelector(".ant-select-disabled");
+      expect(disabledSelect).toBeTruthy();
+    });
+
+    it("should not disable the organization dropdown for admin users", async () => {
+      const { container } = renderWithProviders(
+        <KeyEditView
+          keyData={MOCK_KEY_DATA}
+          onCancel={() => {}}
+          onSubmit={async () => {}}
+          accessToken=""
+          userID=""
+          userRole="Admin"
+          premiumUser={false}
+        />,
+      );
+
+      await waitFor(() => {
+        expect(screen.getByText("Organization")).toBeInTheDocument();
+      });
+
+      const orgFormItem = screen.getByText("Organization").closest(".ant-form-item");
+      const disabledSelect = orgFormItem?.querySelector(".ant-select-disabled");
+      expect(disabledSelect).toBeFalsy();
+    });
+
+    it("should initialize organization from keyData", async () => {
+      const keyWithOrg = {
+        ...MOCK_KEY_DATA,
+        organization_id: "org-1",
+      };
+
+      renderWithProviders(
+        <KeyEditView
+          keyData={keyWithOrg}
+          onCancel={() => {}}
+          onSubmit={async () => {}}
+          accessToken=""
+          userID=""
+          userRole="Admin"
+          premiumUser={false}
+        />,
+      );
+
+      await waitFor(() => {
+        expect(screen.getByText("Engineering")).toBeInTheDocument();
+      });
+    });
+  });
 });

ui/litellm-dashboard/src/components/templates/key_edit_view.tsx

@@ -1,4 +1,5 @@
 import GuardrailSelector from "@/components/guardrails/GuardrailSelector";
+import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations";
 import { useProjects } from "@/app/(dashboard)/hooks/projects/useProjects";
 import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings";
 import PolicySelector from "@/components/policies/PolicySelector";
@@ -13,6 +14,7 @@ import { mapInternalToDisplayNames } from "../callback_info_helpers";
 import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings";
 import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector";
 import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem";
+import OrganizationDropdown from "../common_components/OrganizationDropdown";
 import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils";
 import { KeyResponse } from "../key_team_helpers/key_list";
 import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
@@ -96,10 +98,12 @@ export function KeyEditView({
       ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks)
       : [],
   );
+  const [selectedOrganizationId, setSelectedOrganizationId] = useState<string | null>(keyData.organization_id || null);
   const [autoRotationEnabled, setAutoRotationEnabled] = useState<boolean>(keyData.auto_rotate || false);
   const [rotationInterval, setRotationInterval] = useState<string>(keyData.rotation_interval || "");
   const [neverExpire, setNeverExpire] = useState<boolean>(!keyData.expires);
   const [isKeySaving, setIsKeySaving] = useState(false);
+  const { data: organizations, isLoading: isOrganizationsLoading } = useOrganizations();
   const { data: projects } = useProjects();
   const { data: uiSettingsData } = useUISettings();
   const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui);
@@ -610,6 +614,28 @@ export function KeyEditView({
         />
       </Form.Item>
 
+      <Form.Item
+        label={
+          <span>
+            Organization{" "}
+            <Tooltip title="The organization this key belongs to. Selecting an organization filters the available teams.">
+              <InfoCircleOutlined style={{ marginLeft: "4px" }} />
+            </Tooltip>
+          </span>
+        }
+        name="organization_id"
+      >
+        <OrganizationDropdown
+          organizations={organizations}
+          loading={isOrganizationsLoading}
+          disabled={userRole !== "Admin"}
+          onChange={(orgId) => {
+            setSelectedOrganizationId(orgId || null);
+            form.setFieldValue("team_id", undefined);
+          }}
+        />
+      </Form.Item>
+
       <Form.Item
         label="Team ID"
         name="team_id"
@@ -620,13 +646,29 @@ export function KeyEditView({
           showSearch
           disabled={enableProjectsUI && hasProject}
           style={{ width: "100%" }}
+          onChange={(teamId) => {
+            const selectedTeam = teams?.find((t) => t.team_id === teamId) || null;
+            if (selectedTeam?.organization_id) {
+              setSelectedOrganizationId(selectedTeam.organization_id);
+              form.setFieldValue("organization_id", selectedTeam.organization_id);
+            } else if (!teamId) {
+              setSelectedOrganizationId(null);
+              form.setFieldValue("organization_id", undefined);
+            }
+          }}
           filterOption={(input, option) => {
-            const team = teams?.find((t) => t.team_id === option?.value);
+            const filteredTeams = selectedOrganizationId
+              ? teams?.filter((t) => t.organization_id === selectedOrganizationId)
+              : teams;
+            const team = filteredTeams?.find((t) => t.team_id === option?.value);
             if (!team) return false;
             return team.team_alias?.toLowerCase().includes(input.toLowerCase()) ?? false;
           }}
         >
-          {teams?.map((team) => (
+          {(selectedOrganizationId
+            ? teams?.filter((t) => t.organization_id === selectedOrganizationId)
+            : teams
+          )?.map((team) => (
             <Select.Option key={team.team_id} value={team.team_id}>
               {`${team.team_alias} (${team.team_id})`}
             </Select.Option>