From a3236ba9479e32ee914d43f12082f4423403e3ca Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 9 Oct 2026 00:30:44 -0700 Subject: [PATCH] feat(ui): configure OpenAI workload identity federation from the LLM Credentials and Add Model forms (#45528) * feat(ui): configure OpenAI workload identity federation from the LLM Credentials and Add Model forms * fix(ui): let a federated OpenAI credential omit the service account when the proxy env provides it * fix(ui): clear the federation API Base error when the admin switches the credential to an API key --- .../AddModelForm.integration.test.tsx | 26 +- .../src/components/add_model/AddModelForm.tsx | 5 +- .../add_model/provider_specific_fields.tsx | 15 +- .../CredentialModal.integration.test.tsx | 140 +++++++ .../components/model_add/CredentialModal.tsx | 64 ++-- .../model_add/CredentialsTable.test.tsx | 21 +- .../model_add/CredentialsTableColumns.tsx | 4 +- ...erationFields.tsx => FederationFields.tsx} | 104 ++++-- .../model_add/anthropic_federation.test.ts | 235 +----------- .../model_add/anthropic_federation.ts | 126 +------ .../model_add/credential_federation.test.ts | 345 ++++++++++++++++++ .../model_add/credential_federation.ts | 145 ++++++++ .../model_add/federation_field.test.ts | 19 + .../components/model_add/federation_field.ts | 27 ++ .../model_add/openai_federation.test.ts | 38 ++ .../components/model_add/openai_federation.ts | 46 +++ 16 files changed, 949 insertions(+), 411 deletions(-) rename ui/litellm-dashboard/src/components/model_add/{AnthropicFederationFields.tsx => FederationFields.tsx} (62%) create mode 100644 ui/litellm-dashboard/src/components/model_add/credential_federation.test.ts create mode 100644 ui/litellm-dashboard/src/components/model_add/credential_federation.ts create mode 100644 ui/litellm-dashboard/src/components/model_add/federation_field.test.ts create mode 100644 ui/litellm-dashboard/src/components/model_add/federation_field.ts create mode 100644 ui/litellm-dashboard/src/components/model_add/openai_federation.test.ts create mode 100644 ui/litellm-dashboard/src/components/model_add/openai_federation.ts diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx index 45014d06a3a..e80d0e554e6 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx @@ -464,8 +464,32 @@ describe("AddModelForm", () => { expect(providerSelect).toBeDisabled(); }); + it("saves an OpenAI federated credential and attaches it to the model", async () => { + const user = userEvent.setup(); + const props = await renderAsRole("proxy_admin", Providers.OpenAI); + + await user.click(screen.getByRole("button", { name: "Use workload identity federation" })); + expect(await screen.findByRole("combobox", { name: /^Authentication:/ })).toHaveTextContent( + "Workload identity federation", + ); + fill("Credential Name:", "openai-federated"); + fill(/Service Account ID/, "svc_new"); + await user.click(screen.getByRole("button", { name: "Add Credential" })); + + await waitFor(() => { + expect(credentialCreateCall).toHaveBeenCalledWith("test-access-token", { + credential_name: "openai-federated", + credential_values: { openai_service_account_id: "svc_new" }, + credential_info: { custom_llm_provider: Providers.OpenAI }, + }); + }); + await waitFor(() => { + expect(props.form.getValues("litellm_credential_name")).toBe("openai-federated"); + }); + }); + it("is not offered for a provider without federation support", async () => { - await renderAsRole("proxy_admin", Providers.OpenAI); + await renderAsRole("proxy_admin", Providers.OpenAI_Compatible); expect(screen.queryByRole("button", { name: "Use workload identity federation" })).not.toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index 2b3593e0149..e50750318c7 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -27,7 +27,7 @@ import { import type { Team } from "../key_team_helpers/key_list"; import { type CredentialItem, type ProviderCreateInfo, credentialCreateCall, modelAvailableCall } from "../networking"; import CredentialModal from "../model_add/CredentialModal"; -import { isAnthropicProvider } from "../model_add/anthropic_federation"; +import { federatedProviderOf } from "../model_add/credential_federation"; import { buildCredential, withoutRestrictedFields } from "../model_add/credential_form_helpers"; import { ProviderLogo } from "../molecules/models/ProviderLogo"; import AccessGroupTagsCombobox from "./AccessGroupTagsCombobox"; @@ -100,7 +100,8 @@ const AddModelForm: React.FC = ({ const selectedCredentialName = useWatch({ control: form.control, name: "litellm_credential_name" }); const queryClient = useQueryClient(); const [isFederatedCredentialModalOpen, setIsFederatedCredentialModalOpen] = useState(false); - const canCreateFederatedCredential = isProxyAdminRole(userRole ?? "") && isAnthropicProvider(selectedProvider); + const canCreateFederatedCredential = + isProxyAdminRole(userRole ?? "") && federatedProviderOf(selectedProvider) !== null; const handleCreateFederatedCredential = async (values: Record) => { const credential = buildCredential(values, withoutRestrictedFields(values)); diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx index ab221652771..69f11897009 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx @@ -16,10 +16,12 @@ import { import { ProviderCredentialFieldMetadata } from "../networking"; import { Providers } from "../provider_info_helpers"; import { labelWithHint } from "@/components/shared/form/LabelWithHint"; +import type { ProviderFieldValidators } from "../model_add/credential_federation"; interface ProviderSpecificFieldsProps { selectedProvider: string | null; hiddenFieldKeys?: readonly string[]; + fieldValidators?: ProviderFieldValidators; } const readTextFile = (file: File, onLoaded: (contents: string) => void) => { @@ -81,7 +83,11 @@ const mapFieldMetadataToUiField = (field: ProviderCredentialFieldMetadata): Prov const providerFieldsByDisplayName: Record = {}; -const ProviderSpecificFields: React.FC = ({ selectedProvider, hiddenFieldKeys }) => { +const ProviderSpecificFields: React.FC = ({ + selectedProvider, + hiddenFieldKeys, + fieldValidators, +}) => { const selectedProviderEnum = Providers[selectedProvider as keyof typeof Providers] as Providers; const form = useFormContext(); const credentialsFileRef = React.useRef(null); @@ -297,7 +303,12 @@ const ProviderSpecificFields: React.FC = ({ selecte label={field.tooltip ? labelWithHint(field.label, field.tooltip) : field.label} name={field.key} required={field.required} - rules={field.required ? { validate: { required: requiredRule("Required") } } : undefined} + rules={{ + validate: { + ...(field.required ? { required: requiredRule("Required") } : {}), + ...(fieldValidators?.[field.key] ? { provider: fieldValidators[field.key] } : {}), + }, + }} className={field.key === "vertex_credentials" ? "mb-0" : "mb-4"} > {(control) => renderFieldControl(field, control)} diff --git a/ui/litellm-dashboard/src/components/model_add/CredentialModal.integration.test.tsx b/ui/litellm-dashboard/src/components/model_add/CredentialModal.integration.test.tsx index d1b54219dd3..32ed7762e70 100644 --- a/ui/litellm-dashboard/src/components/model_add/CredentialModal.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/model_add/CredentialModal.integration.test.tsx @@ -442,3 +442,143 @@ describe("CredentialModal with Anthropic workload identity federation", () => { ); }); }); + +const openAIFederatedCredential: CredentialItem = { + credential_name: "openai-federated", + credential_values: { + api_base: "https://api.openai.com/v1", + openai_identity_provider_id: "idp_stored", + openai_service_account_id: "svc_stored", + openai_identity_token_file: "/var****", + }, + credential_info: { custom_llm_provider: "openai" }, +}; + +describe("CredentialModal with OpenAI workload identity federation", () => { + it("creates a federated OpenAI credential and never sends an API key", async () => { + const user = userEvent.setup(); + const onSubmit = renderModal({ initialProvider: "OpenAI" }); + await screen.findByLabelText("OpenAI API Key"); + fill("OpenAI API Key", "sk-proj-typed-before-switching"); + + await chooseOption(user, /^Authentication:/, "Workload identity federation"); + + expect(screen.queryByLabelText("OpenAI API Key")).not.toBeInTheDocument(); + expect(screen.getByText(/only when OPENAI_API_KEY is unset/)).toBeInTheDocument(); + fill("Credential Name:", "openai-federated"); + fill(/Identity Provider ID/, "idp_new"); + fill(/Service Account ID/, " svc_new "); + fill(/Identity Token File/, "/var/run/secrets/kubernetes.io/serviceaccount/token"); + await user.click(screen.getByRole("button", { name: "Add Credential" })); + + const expectedPayload = { + credential_name: "openai-federated", + custom_llm_provider: "OpenAI", + openai_identity_provider_id: "idp_new", + openai_service_account_id: "svc_new", + openai_identity_token_file: "/var/run/secrets/kubernetes.io/serviceaccount/token", + }; + expect(onSubmit).toHaveBeenCalledWith(expectedPayload, []); + }); + + it("saves a stored federated OpenAI credential untouched when its service account comes from the proxy environment", async () => { + const user = userEvent.setup(); + const { openai_service_account_id: _, ...valuesWithoutServiceAccount } = + openAIFederatedCredential.credential_values; + const onSubmit = renderModal({ + mode: "edit", + existingCredential: { ...openAIFederatedCredential, credential_values: valuesWithoutServiceAccount }, + }); + await screen.findByLabelText(/Service Account ID/); + + await user.click(screen.getByRole("button", { name: "Update Credential" })); + + expect(onSubmit).toHaveBeenCalledWith({ credential_name: "openai-federated", custom_llm_provider: "openai" }, []); + }); + + it("refuses a base URL the proxy would not federate with, and drops that check once the admin picks an API key", async () => { + const user = userEvent.setup(); + const onSubmit = renderModal({ initialProvider: "OpenAI", initialAuthMethod: "federation" }); + await screen.findByLabelText("API Base"); + fill("Credential Name:", "openai-gateway"); + fill(/Service Account ID/, "svc_new"); + fill("API Base", "https://gateway.example.com/v1"); + + await user.click(screen.getByRole("button", { name: "Add Credential" })); + + expect(await screen.findByText(/only reaches the OpenAI API/)).toBeInTheDocument(); + expect(onSubmit).not.toHaveBeenCalled(); + + await chooseOption(user, /^Authentication:/, "API key"); + await screen.findByLabelText("OpenAI API Key"); + expect(screen.queryByText(/only reaches the OpenAI API/)).not.toBeInTheDocument(); + fill("OpenAI API Key", "sk-proj-new"); + await user.click(screen.getByRole("button", { name: "Add Credential" })); + + const expectedPayload = { + credential_name: "openai-gateway", + custom_llm_provider: "OpenAI", + api_base: "https://gateway.example.com/v1", + api_key: "sk-proj-new", + }; + expect(onSubmit).toHaveBeenCalledWith(expectedPayload, []); + }); + + it("shows a stored federated OpenAI credential and writes nothing when it is saved untouched", async () => { + const user = userEvent.setup(); + const onSubmit = renderModal({ mode: "edit", existingCredential: openAIFederatedCredential }); + + expect(await screen.findByRole("combobox", { name: /^Authentication:/ })).toHaveTextContent( + "Workload identity federation", + ); + expect(screen.getByLabelText(/Service Account ID/)).toHaveValue("svc_stored"); + expect(screen.getByLabelText(/Identity Token File/)).toHaveValue("/var****"); + expect(screen.queryByLabelText("OpenAI API Key")).not.toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Update Credential" })); + + expect(onSubmit).toHaveBeenCalledWith({ credential_name: "openai-federated", custom_llm_provider: "openai" }, []); + }); + + it("deletes the stored OpenAI federation values when the admin switches the credential to an API key", async () => { + const user = userEvent.setup(); + const onSubmit = renderModal({ mode: "edit", existingCredential: openAIFederatedCredential }); + await screen.findByLabelText(/Service Account ID/); + + await chooseOption(user, /^Authentication:/, "API key"); + await screen.findByLabelText("OpenAI API Key"); + fill("OpenAI API Key", "sk-proj-replacement"); + await user.click(screen.getByRole("button", { name: "Update Credential" })); + + const [values, valuesToDelete] = onSubmit.mock.calls[0]; + expect(values).toEqual({ + credential_name: "openai-federated", + custom_llm_provider: "openai", + api_key: "sk-proj-replacement", + }); + expect([...valuesToDelete].sort()).toEqual([ + "openai_identity_provider_id", + "openai_identity_token_file", + "openai_service_account_id", + ]); + }); + + it("restores the stored federation settings when the admin returns to the credential's own provider", async () => { + const user = userEvent.setup(); + const onSubmit = renderModal({ mode: "edit", existingCredential: federatedCredential }); + await screen.findByLabelText(/Federation Rule ID/); + + await chooseProvider(user, "OpenAI"); + expect(await screen.findByRole("combobox", { name: /^Authentication:/ })).toHaveTextContent("API key"); + await chooseProvider(user, "Anthropic"); + + expect(await screen.findByRole("combobox", { name: /Identity Source/ })).toHaveTextContent( + "Keycloak client credentials", + ); + await user.click(screen.getByRole("button", { name: "Update Credential" })); + expect(onSubmit).toHaveBeenCalledWith( + { credential_name: "anthropic-federated", custom_llm_provider: "Anthropic" }, + [], + ); + }); +}); diff --git a/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx b/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx index 61afd5eea99..65778e72f2d 100644 --- a/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx +++ b/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx @@ -19,18 +19,19 @@ import { Providers } from "../provider_info_helpers"; import { Logo } from "@/components/molecules/logo/Logo"; import { resetCredentialFormOnProviderChange, withoutRestrictedFields } from "./credential_form_helpers"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; -import AnthropicFederationFields from "./AnthropicFederationFields"; +import FederationFields from "./FederationFields"; +import { DEFAULT_IDENTITY_SOURCE, inferIdentitySource, type IdentitySourceId } from "./anthropic_federation"; import { buildCreateCredentialValues, buildCredentialPatch, buildProviderChangePatch, + federatedProviderOf, inferAuthMethod, - inferIdentitySource, - isAnthropicProvider, isFederatedCredential, - type AnthropicAuthMethod, - type IdentitySourceId, -} from "./anthropic_federation"; + providerFieldValidators, + selectionFor, + type AuthMethod, +} from "./credential_federation"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; const providerOptions: SearchSelectOption[] = Object.entries(Providers).map(([providerEnum, providerDisplayName]) => ({ @@ -39,11 +40,11 @@ const providerOptions: SearchSelectOption[] = Object.entries(Providers).map(([pr icon: , })); -const AUTH_METHOD_SELECT_ID = "anthropic_auth_method"; +const AUTH_METHOD_SELECT_ID = "credential_auth_method"; const API_KEY_FIELDS: readonly string[] = ["api_key"]; const NO_HIDDEN_FIELDS: readonly string[] = []; -const authMethodItems: { value: AnthropicAuthMethod; label: string }[] = [ +const authMethodItems: { value: AuthMethod; label: string }[] = [ { value: "api_key", label: "API key" }, { value: "federation", label: "Workload identity federation" }, ]; @@ -55,7 +56,7 @@ interface CredentialModalProps { mode: "add" | "edit"; existingCredential?: CredentialItem | null; initialProvider?: string | null; - initialAuthMethod?: AnthropicAuthMethod; + initialAuthMethod?: AuthMethod; providerLocked?: boolean; } @@ -111,18 +112,16 @@ export default function CredentialModal({ ); const storedProvider = existingCredential?.credential_info.custom_llm_provider ?? null; const storedValues: Record = existingCredential?.credential_values ?? {}; - const storedSelection = { - authMethod: inferAuthMethod(storedValues), - identitySource: inferIdentitySource(storedValues), - }; - const [authMethod, setAuthMethod] = useState( - existingCredential ? storedSelection.authMethod : initialAuthMethod ?? "api_key", + const storedAuthMethod = inferAuthMethod(storedValues); + const storedIdentitySource = isFederatedCredential(storedValues) + ? inferIdentitySource(storedValues) + : DEFAULT_IDENTITY_SOURCE; + const storedSelection = selectionFor(federatedProviderOf(storedProvider), storedAuthMethod, storedIdentitySource); + const [authMethod, setAuthMethod] = useState( + existingCredential ? storedAuthMethod : initialAuthMethod ?? "api_key", ); - const [identitySource, setIdentitySource] = useState( - isFederatedCredential(storedValues) ? storedSelection.identitySource : "token_file", - ); - const isAnthropic = isAnthropicProvider(selectedProvider); - const selection = { authMethod: isAnthropic ? authMethod : ("api_key" as const), identitySource }; + const [identitySource, setIdentitySource] = useState(storedIdentitySource); + const selection = selectionFor(federatedProviderOf(selectedProvider), authMethod, identitySource); const initialValues = initialFormValues(existingCredential, initialProvider); @@ -135,6 +134,18 @@ export default function CredentialModal({ setFieldValue: (field: string, value: unknown) => form.setValue(field, value), }); + const changeProvider = (provider: string | null) => { + const backToStored = isEdit && sameProvider(provider, storedProvider); + setAuthMethod(backToStored ? storedAuthMethod : "api_key"); + setIdentitySource(backToStored ? storedIdentitySource : DEFAULT_IDENTITY_SOURCE); + resetCredentialFormOnProviderChange(formAdapterFor(provider), provider, setSelectedProvider); + }; + + const changeAuthMethod = (method: AuthMethod) => { + form.clearErrors(Object.keys(providerFieldValidators(selection))); + setAuthMethod(method); + }; + const handleSubmit = async () => { const isValid = await form.trigger(registry.mountedNames() as string[]); if (!isValid) { @@ -234,24 +245,24 @@ export default function CredentialModal({ disabled={providerLocked} onValueChange={(value) => { control.onChange(value); - resetCredentialFormOnProviderChange(formAdapterFor(value), value, setSelectedProvider); + changeProvider(value); }} /> )} - {isAnthropic && ( + {federatedProviderOf(selectedProvider) !== null && (