litellm/ui/litellm-dashboard/src/components/UIAccessControlForm.tsx

import React, { useEffect, useState } from "react";
import { Form, Button as Button2, Select, message } from "antd";
import { Text, TextInput } from "@tremor/react";
import { getSSOSettings, updateSSOSettings } from "./networking";

interface UIAccessControlFormProps {
  accessToken: string | null;
  onSuccess: () => void;
}

// Separate UI Access Control Form Component
const UIAccessControlForm: React.FC<UIAccessControlFormProps> = ({ accessToken, onSuccess }) => {
  const [form] = Form.useForm();
  const [loading, setLoading] = useState(false);

  // Load existing UI access control settings
  useEffect(() => {
    const loadUIAccessSettings = async () => {
      if (accessToken) {
        try {
          const ssoData = await getSSOSettings(accessToken);
          if (ssoData && ssoData.values) {
            // Handle nested ui_access_mode structure
            const uiAccessMode = ssoData.values.ui_access_mode;
            let formValues = {};
            
            if (uiAccessMode && typeof uiAccessMode === 'object') {
              formValues = {
                ui_access_mode_type: uiAccessMode.type,
                restricted_sso_group: uiAccessMode.restricted_sso_group,
                sso_group_jwt_field: uiAccessMode.sso_group_jwt_field,
              };
            } else if (typeof uiAccessMode === 'string') {
              // Handle legacy flat structure
              formValues = {
                ui_access_mode_type: uiAccessMode,
                restricted_sso_group: ssoData.values.restricted_sso_group,
                sso_group_jwt_field: ssoData.values.team_ids_jwt_field || ssoData.values.sso_group_jwt_field,
              };
            }
            
            form.setFieldsValue(formValues);
          }
        } catch (error) {
          console.error("Failed to load UI access settings:", error);
        }
      }
    };

    loadUIAccessSettings();
  }, [accessToken, form]);

  const handleUIAccessSubmit = async (formValues: Record<string, any>) => {
    if (!accessToken) {
      message.error("No access token available");
      return;
    }

    setLoading(true);
    try {
      // Transform form data to match API expected structure
      const apiPayload = {
        ui_access_mode: {
          type: formValues.ui_access_mode_type,
          restricted_sso_group: formValues.restricted_sso_group,
          sso_group_jwt_field: formValues.sso_group_jwt_field,
        }
      };

      await updateSSOSettings(accessToken, apiPayload);
      onSuccess();
    } catch (error) {
      console.error("Failed to save UI access settings:", error);
      message.error("Failed to save UI access settings");
    } finally {
      setLoading(false);
    }
  };

  return (
    <div style={{ padding: '16px' }}>
      <div style={{ marginBottom: '16px' }}>
        <Text style={{ fontSize: '14px', color: '#6b7280' }}>
          Configure who can access the UI interface and how group information is extracted from JWT tokens.
        </Text>
      </div>
      
      <Form
        form={form}
        onFinish={handleUIAccessSubmit}
        layout="vertical"
      >
        <Form.Item
          label="UI Access Mode"
          name="ui_access_mode_type"
          tooltip="Controls who can access the UI interface"
        >
          <Select placeholder="Select access mode">
            <Select.Option value="all_authenticated_users">All Authenticated Users</Select.Option>
            <Select.Option value="restricted_sso_group">Restricted SSO Group</Select.Option>
          </Select>
        </Form.Item>

        <Form.Item
          noStyle
          shouldUpdate={(prevValues, currentValues) => prevValues.ui_access_mode_type !== currentValues.ui_access_mode_type}
        >
          {({ getFieldValue }) => {
            const uiAccessModeType = getFieldValue('ui_access_mode_type');
            return uiAccessModeType === 'restricted_sso_group' ? (
              <Form.Item
                label="Restricted SSO Group"
                name="restricted_sso_group"
                rules={[{ required: true, message: "Please enter the restricted SSO group" }]}
              >
                <TextInput placeholder="ui-access-group" />
              </Form.Item>
            ) : null;
          }}
        </Form.Item>

        <Form.Item
          label="SSO Group JWT Field"
          name="sso_group_jwt_field"
          tooltip="JWT field name that contains team/group information. Use dot notation to access nested fields."
        >
          <TextInput placeholder="groups" />
        </Form.Item>

        <div style={{ textAlign: "right", marginTop: "16px" }}>
          <Button2 
            type="primary" 
            htmlType="submit" 
            loading={loading}
            style={{ 
              backgroundColor: '#6366f1', 
              borderColor: '#6366f1' 
            }}
          >
            Update UI Access Control
          </Button2>
        </div>
      </Form>
    </div>
  );
};

export default UIAccessControlForm;