From ce8179b2c0f130eb8ab6c025be56f1f1d7ee827b Mon Sep 17 00:00:00 2001
From: mateo-berri <edgetonmas@gmail.com>
Date: Thu, 3 Sep 2026 03:04:33 -0700
Subject: [PATCH] fix(ui): keep the key edit team filter in sync with the
 organization field

---
 .../templates/key_edit_view.test.tsx          | 34 +++++++++++++++++++
 .../components/templates/key_edit_view.tsx    |  4 +--
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx
index dfa3e8b2bca..feb777f32a9 100644
--- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx
+++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx
@@ -1468,6 +1468,40 @@ describe("KeyEditView", () => {
       expect(screen.queryByRole("option", { name: /Beta/ })).not.toBeInTheDocument();
     });
 
+    it("should re-scope teams when the edited key switches to another organization", async () => {
+      const sharedProps = {
+        teams: [
+          { team_id: "team-a", team_alias: "Alpha", organization_id: "org-1" },
+          { team_id: "team-b", team_alias: "Beta", organization_id: "org-2" },
+        ],
+        onCancel: () => {},
+        onSubmit: async () => {},
+        accessToken: "",
+        userID: "",
+        userRole: "Admin",
+        premiumUser: false,
+      };
+
+      const { rerender } = renderWithProviders(
+        <KeyEditView keyData={{ ...MOCK_KEY_DATA, organization_id: null, org_id: "org-1" }} {...sharedProps} />,
+      );
+
+      await waitFor(() => {
+        expect(screen.getByLabelText("Organization")).toHaveValue("Engineering");
+      });
+
+      rerender(<KeyEditView keyData={{ ...MOCK_KEY_DATA, organization_id: null, org_id: "org-2" }} {...sharedProps} />);
+
+      await waitFor(() => {
+        expect(screen.getByLabelText("Organization")).toHaveValue("Sales");
+      });
+
+      await userEvent.click(screen.getByLabelText("Team ID"));
+
+      expect(await screen.findByRole("option", { name: /Beta/ })).toBeInTheDocument();
+      expect(screen.queryByRole("option", { name: /Alpha/ })).not.toBeInTheDocument();
+    });
+
     it("should initialize organization from keyData", async () => {
       const keyWithOrg = {
         ...MOCK_KEY_DATA,
diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx
index e925eac5380..3242aee49bc 100644
--- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx
+++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx
@@ -28,7 +28,6 @@ import { routerSettingsEditorValue, routerSettingsUpdate } from "../common_compo
 import { estimateTooltips, withNormalizedEstimates } from "./estimatedOutputTokens";
 import {
   currentValuePlaceholder,
-  keyOrganizationId,
   keyTypeFromRoutes,
   modelSentinelOptions,
   parseAllowedRoutes,
@@ -106,7 +105,6 @@ export function KeyEditView({
       ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks)
       : [],
   );
-  const [selectedOrganizationId, setSelectedOrganizationId] = useState<string | null>(keyOrganizationId(keyData));
   const [autoRotationEnabled, setAutoRotationEnabled] = useState<boolean>(keyData.auto_rotate || false);
   const [rotationInterval, setRotationInterval] = useState<string>(keyData.rotation_interval || "");
   const [neverExpire, setNeverExpire] = useState<boolean>(!keyData.expires);
@@ -135,6 +133,7 @@ export function KeyEditView({
     return project?.project_alias ? `${project.project_alias} (${keyData.project_id})` : keyData.project_id;
   })();
 
+  const selectedOrganizationId = form.watch("organization_id") ?? null;
   const allowedRoutesValue = form.watch("allowed_routes");
   const selectedModels = (form.watch("models") as string[] | undefined) ?? [];
   const allowedRoutes = parseAllowedRoutes(allowedRoutesValue);
@@ -306,7 +305,6 @@ export function KeyEditView({
 
   const handleOrganizationChange = (setField: (value: string | undefined) => void, orgId: string | undefined) => {
     setField(orgId);
-    setSelectedOrganizationId(orgId || null);
     form.setValue("team_id", undefined);
   };