[Feature] UI - SSO: Add Team Mappings

ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts

@@ -28,6 +28,7 @@ export interface SSOSettingsValues {
   user_email: string | null;
   ui_access_mode: string | null;
   role_mappings: RoleMappings;
+  team_mappings: TeamMappings;
 }
 
 export interface RoleMappings {
@@ -39,6 +40,10 @@ export interface RoleMappings {
   };
 }
 
+export interface TeamMappings {
+  team_ids_jwt_field: string;
+}
+
 export interface SSOSettingsResponse {
   values: SSOSettingsValues;
   field_schema: SSOFieldSchema;

ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.test.tsx

@@ -151,6 +151,117 @@ describe("BaseSSOSettingsForm", () => {
       expect(screen.getByText("Default Role")).toBeInTheDocument();
     });
   });
+
+  it("should show team mappings checkbox for okta provider", async () => {
+    const TestWrapper = () => {
+      const [form] = Form.useForm();
+      const handleSubmit = vi.fn();
+
+      return <BaseSSOSettingsForm form={form} onFormSubmit={handleSubmit} />;
+    };
+
+    renderWithProviders(<TestWrapper />);
+
+    const providerSelect = screen.getByLabelText("SSO Provider");
+    await act(async () => {
+      fireEvent.mouseDown(providerSelect);
+    });
+
+    await waitFor(() => {
+      const oktaOption = screen.getByText(/okta/i);
+      fireEvent.click(oktaOption);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Use Team Mappings")).toBeInTheDocument();
+    });
+  });
+
+  it("should show team mappings checkbox for generic provider", async () => {
+    const TestWrapper = () => {
+      const [form] = Form.useForm();
+      const handleSubmit = vi.fn();
+
+      return <BaseSSOSettingsForm form={form} onFormSubmit={handleSubmit} />;
+    };
+
+    renderWithProviders(<TestWrapper />);
+
+    const providerSelect = screen.getByLabelText("SSO Provider");
+    await act(async () => {
+      fireEvent.mouseDown(providerSelect);
+    });
+
+    await waitFor(() => {
+      const genericOption = screen.getByText(/generic sso/i);
+      fireEvent.click(genericOption);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Use Team Mappings")).toBeInTheDocument();
+    });
+  });
+
+  it("should show team IDs JWT field when use_team_mappings is checked for okta provider", async () => {
+    const TestWrapper = () => {
+      const [form] = Form.useForm();
+      const handleSubmit = vi.fn();
+
+      return <BaseSSOSettingsForm form={form} onFormSubmit={handleSubmit} />;
+    };
+
+    renderWithProviders(<TestWrapper />);
+
+    const providerSelect = screen.getByLabelText("SSO Provider");
+    await act(async () => {
+      fireEvent.mouseDown(providerSelect);
+    });
+
+    await waitFor(() => {
+      const oktaOption = screen.getByText(/okta/i);
+      fireEvent.click(oktaOption);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Use Team Mappings")).toBeInTheDocument();
+    });
+
+    const checkbox = screen.getByLabelText("Use Team Mappings");
+    await act(async () => {
+      fireEvent.click(checkbox);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Team IDs JWT Field")).toBeInTheDocument();
+    });
+  });
+
+  it("should not show team mappings checkbox for google provider", async () => {
+    const TestWrapper = () => {
+      const [form] = Form.useForm();
+      const handleSubmit = vi.fn();
+
+      return <BaseSSOSettingsForm form={form} onFormSubmit={handleSubmit} />;
+    };
+
+    renderWithProviders(<TestWrapper />);
+
+    const providerSelect = screen.getByLabelText("SSO Provider");
+    await act(async () => {
+      fireEvent.mouseDown(providerSelect);
+    });
+
+    await waitFor(() => {
+      const googleOption = screen.getByText(/google sso/i);
+      fireEvent.click(googleOption);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Google Client ID")).toBeInTheDocument();
+    });
+
+    expect(screen.queryByText("Use Team Mappings")).not.toBeInTheDocument();
+  });
 });
 
 describe("renderProviderFields", () => {

ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx

@@ -251,6 +251,43 @@ const BaseSSOSettingsForm: React.FC<BaseSSOSettingsFormProps> = ({ form, onFormS
             ) : null;
           }}
         </Form.Item>
+
+        <Form.Item
+          noStyle
+          shouldUpdate={(prevValues, currentValues) => prevValues.sso_provider !== currentValues.sso_provider}
+        >
+          {({ getFieldValue }) => {
+            const provider = getFieldValue("sso_provider");
+            return provider === "okta" || provider === "generic" ? (
+              <Form.Item label="Use Team Mappings" name="use_team_mappings" valuePropName="checked">
+                <Checkbox />
+              </Form.Item>
+            ) : null;
+          }}
+        </Form.Item>
+
+        <Form.Item
+          noStyle
+          shouldUpdate={(prevValues, currentValues) =>
+            prevValues.use_team_mappings !== currentValues.use_team_mappings ||
+            prevValues.sso_provider !== currentValues.sso_provider
+          }
+        >
+          {({ getFieldValue }) => {
+            const useTeamMappings = getFieldValue("use_team_mappings");
+            const provider = getFieldValue("sso_provider");
+            const supportsTeamMappings = provider === "okta" || provider === "generic";
+            return useTeamMappings && supportsTeamMappings ? (
+              <Form.Item
+                label="Team IDs JWT Field"
+                name="team_ids_jwt_field"
+                rules={[{ required: true, message: "Please enter the team IDs JWT field" }]}
+              >
+                <TextInput />
+              </Form.Item>
+            ) : null;
+          }}
+        </Form.Item>
       </Form>
     </div>
   );

ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.tsx

@@ -33,6 +33,7 @@ const DeleteSSOSettingsModal: React.FC<DeleteSSOSettingsModalProps> = ({ isVisib
       user_email: null,
       sso_provider: null,
       role_mappings: null,
+      team_mappings: null,
     };
 
     await editSSOSettings(clearSettings, {

ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx

@@ -105,6 +105,14 @@ const createRoleMappingsSSOData = (overrides: Record<string, any> = {}) =>
     ...overrides,
   });
 
+const createTeamMappingsSSOData = (overrides: Record<string, any> = {}) =>
+  createGenericSSOData({
+    team_mappings: {
+      team_ids_jwt_field: overrides.team_ids_jwt_field || "teams",
+    },
+    ...overrides,
+  });
+
 // Mock utilities
 const createMockHooks = (): {
   useSSOSettings: SSOSettingsHookReturn;
@@ -577,6 +585,104 @@ describe("EditSSOSettingsModal", () => {
         });
       });
     });
+  });
+
+  describe("Team Mappings", () => {
+    it("processes team mappings when team_mappings exists", async () => {
+      const ssoData = createTeamMappingsSSOData();
+
+      setupMocks({
+        useSSOSettings: { data: ssoData, isLoading: false, error: null },
+      });
+
+      renderComponent();
+
+      await waitFor(() => {
+        expect(mockForm.setFieldsValue).toHaveBeenCalledWith({
+          sso_provider: SSO_PROVIDERS.GENERIC,
+          ...ssoData.values,
+          use_team_mappings: true,
+          team_ids_jwt_field: "teams",
+        });
+      });
+    });
+
+    it("handles team mappings with custom JWT field name", async () => {
+      const ssoData = createTeamMappingsSSOData({
+        team_ids_jwt_field: "custom_teams_field",
+      });
+
+      setupMocks({
+        useSSOSettings: { data: ssoData, isLoading: false, error: null },
+      });
+
+      renderComponent();
+
+      await waitFor(() => {
+        expect(mockForm.setFieldsValue).toHaveBeenCalledWith({
+          sso_provider: SSO_PROVIDERS.GENERIC,
+          ...ssoData.values,
+          use_team_mappings: true,
+          team_ids_jwt_field: "custom_teams_field",
+        });
+      });
+    });
+
+    it("handles team mappings and role mappings together", async () => {
+      const ssoData = createGenericSSOData({
+        role_mappings: {
+          group_claim: "groups",
+          default_role: "internal_user",
+          roles: {
+            proxy_admin: ["admin-group"],
+            proxy_admin_viewer: [],
+            internal_user: [],
+            internal_user_viewer: [],
+          },
+        },
+        team_mappings: {
+          team_ids_jwt_field: "teams",
+        },
+      });
+
+      setupMocks({
+        useSSOSettings: { data: ssoData, isLoading: false, error: null },
+      });
+
+      renderComponent();
+
+      await waitFor(() => {
+        expect(mockForm.setFieldsValue).toHaveBeenCalledWith({
+          sso_provider: SSO_PROVIDERS.GENERIC,
+          ...ssoData.values,
+          use_role_mappings: true,
+          group_claim: "groups",
+          default_role: "internal_user",
+          proxy_admin_teams: "admin-group",
+          admin_viewer_teams: "",
+          internal_user_teams: "",
+          internal_viewer_teams: "",
+          use_team_mappings: true,
+          team_ids_jwt_field: "teams",
+        });
+      });
+    });
+
+    it("does not set team mapping fields when team_mappings is not present", async () => {
+      const ssoData = createGenericSSOData();
+
+      setupMocks({
+        useSSOSettings: { data: ssoData, isLoading: false, error: null },
+      });
+
+      renderComponent();
+
+      await waitFor(() => {
+        const callArgs = mockForm.setFieldsValue.mock.calls[0][0];
+        expect(callArgs.use_team_mappings).toBeUndefined();
+        expect(callArgs.team_ids_jwt_field).toBeUndefined();
+      });
+    });
 
     it("handles provider detection with partial SSO data", async () => {
       const ssoData = createSSOData({

ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx

@@ -68,11 +68,22 @@ const EditSSOSettingsModal: React.FC<EditSSOSettingsModalProps> = ({ isVisible,
         };
       }
 
+      // Extract team mappings if they exist
+      let teamMappingFields = {};
+      if (ssoData.values.team_mappings) {
+        const teamMappings = ssoData.values.team_mappings;
+        teamMappingFields = {
+          use_team_mappings: true,
+          team_ids_jwt_field: teamMappings.team_ids_jwt_field,
+        };
+      }
+
       // Set form values with existing data (excluding UI access control fields)
       const formValues = {
         sso_provider: selectedProvider,
         ...ssoData.values,
         ...roleMappingFields,
+        ...teamMappingFields,
       };
 
       console.log("Setting form values:", formValues); // Debug log

ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import { useSSOSettings, type SSOSettingsValues } from "@/app/(dashboard)/hooks/sso/useSSOSettings";
-import { Button, Card, Descriptions, Space, Typography } from "antd";
+import { Button, Card, Descriptions, Space, Tag, Typography } from "antd";
 import { Edit, Shield, Trash2 } from "lucide-react";
 import { useState } from "react";
 import { ssoProviderDisplayNames, ssoProviderLogoMap } from "./constants";
@@ -28,6 +28,7 @@ export default function SSOSettings() {
 
   const selectedProvider = ssoSettings?.values ? detectSSOProvider(ssoSettings.values) : null;
   const isRoleMappingsEnabled = Boolean(ssoSettings?.values.role_mappings);
+  const isTeamMappingsEnabled = Boolean(ssoSettings?.values.team_mappings);
 
   const renderEndpointValue = (value?: string | null) => (
     <Text className="font-mono text-gray-600 text-sm" copyable={!!value}>
@@ -38,6 +39,15 @@ export default function SSOSettings() {
   const renderSimpleValue = (value?: string | null) =>
     value ? value : <span className="text-gray-400 italic">Not configured</span>;
 
+  const renderTeamMappingsField = (values: SSOSettingsValues) => {
+    if (!values.team_mappings?.team_ids_jwt_field) {
+      return <span className="text-gray-400 italic">Not configured</span>;
+    }
+    return (
+      <Tag>{values.team_mappings.team_ids_jwt_field}</Tag>
+    );
+  };
+
   const descriptionsConfig = {
     column: {
       xxl: 1,
@@ -103,6 +113,10 @@ export default function SSOSettings() {
           render: (values: SSOSettingsValues) => renderEndpointValue(values.generic_userinfo_endpoint),
         },
         { label: "Proxy Base URL", render: (values: SSOSettingsValues) => renderSimpleValue(values.proxy_base_url) },
+        isTeamMappingsEnabled ? {
+          label: "Team IDs JWT Field",
+          render: (values: SSOSettingsValues) => renderTeamMappingsField(values),
+        } : null,
       ],
     },
     generic: {
@@ -129,6 +143,10 @@ export default function SSOSettings() {
           render: (values: SSOSettingsValues) => renderEndpointValue(values.generic_userinfo_endpoint),
         },
         { label: "Proxy Base URL", render: (values: SSOSettingsValues) => renderSimpleValue(values.proxy_base_url) },
+        isTeamMappingsEnabled ? {
+          label: "Team IDs JWT Field",
+          render: (values: SSOSettingsValues) => renderTeamMappingsField(values),
+        } : null,
       ],
     },
   };
@@ -155,7 +173,7 @@ export default function SSOSettings() {
             <span>{config.providerText}</span>
           </div>
         </Descriptions.Item>
-        {config.fields.map((field, index) => (
+        {config.fields.map((field, index) => field && (
           <Descriptions.Item key={index} label={field.label}>
             {field.render(values)}
           </Descriptions.Item>