Merge pull request #34452 from BerriAI/litellm_lit4766_key_edit_mcp_toolsets

fix(ui): keep a key's MCP toolsets when saving an edit
This commit is contained in:
tin-berri 2026-07-24 15:38:33 -07:00 • committed by GitHub
commit 57ed2ed718
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 96 additions and 3 deletions

View file

@ -94,6 +94,7 @@ export interface KeyResponse {
object_permission_id: string;
mcp_servers: string[];
mcp_access_groups?: string[];
mcp_toolsets?: string[] | null;
mcp_tool_permissions?: Record<string, string[]>;
vector_stores: string[];
agents?: string[];

View file

@ -9,7 +9,7 @@ interface ObjectPermission {
mcp_servers: string[];
mcp_access_groups?: string[];
mcp_tool_permissions?: Record<string, string[]>;
mcp_toolsets?: string[];
mcp_toolsets?: string[] | null;
vector_stores: string[];
agents?: string[];
agent_access_groups?: string[];

View file

@ -443,6 +443,36 @@ describe("CreateKey", () => {
});
});
it("should include mcp_toolsets in keyCreateCall payload when only toolsets are selected", async () => {
renderWithProviders(<CreateKey {...defaultProps} />);
act(() => {
fireEvent.click(screen.getByRole("button", { name: /create new key/i }));
});
await waitFor(() => {
expect(screen.getByRole("button", { name: /create key/i })).toBeInTheDocument();
});
act(() => {
formMock.setFieldValue("key_alias", "Test Key");
formMock.setFieldValue("allowed_mcp_servers_and_groups", {
servers: [],
accessGroups: [],
toolsets: ["ts-1"],
});
});
act(() => {
fireEvent.click(screen.getByRole("button", { name: /create key/i }));
});
await waitFor(() => {
expect(mockKeyCreateCall).toHaveBeenCalled();
});
expect(mockKeyCreateCall.mock.calls[0][2].object_permission?.mcp_toolsets).toEqual(["ts-1"]);
});
it("should prefill models when provided without team_id", async () => {
renderWithProviders(
<CreateKey

View file

@ -468,18 +468,22 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
if (
formValues.allowed_mcp_servers_and_groups &&
(formValues.allowed_mcp_servers_and_groups.servers?.length > 0 ||
formValues.allowed_mcp_servers_and_groups.accessGroups?.length > 0)
formValues.allowed_mcp_servers_and_groups.accessGroups?.length > 0 ||
formValues.allowed_mcp_servers_and_groups.toolsets?.length > 0)
) {
if (!formValues.object_permission) {
formValues.object_permission = {};
}
const { servers, accessGroups } = formValues.allowed_mcp_servers_and_groups;
const { servers, accessGroups, toolsets } = formValues.allowed_mcp_servers_and_groups;
if (servers && servers.length > 0) {
formValues.object_permission.mcp_servers = servers;
}
if (accessGroups && accessGroups.length > 0) {
formValues.object_permission.mcp_access_groups = accessGroups;
}
if (toolsets && toolsets.length > 0) {
formValues.object_permission.mcp_toolsets = toolsets;
}
// Remove the original field as it's now part of object_permission
delete formValues.allowed_mcp_servers_and_groups;
}

View file

@ -453,6 +453,28 @@ describe("KeyInfoView handleKeyUpdate guardrails guard", () => {
});
});
describe("KeyInfoView handleKeyUpdate mcp_toolsets", () => {
it("should forward the toolsets the edit form supplies into object_permission", async () => {
renderView(true);
fireEvent.click(screen.getByText("Settings"));
fireEvent.click(screen.getByText("Edit Settings"));
(globalThis as any).__TEST_FORM_VALUES = {
token: "tok_123",
max_budget: 40000,
mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: ["ts-1"] },
};
fireEvent.click(screen.getByText("Mock Submit"));
await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled());
const [, sentPayload] = keyUpdateCallMock.mock.calls[0];
expect(sentPayload.object_permission.mcp_toolsets).toEqual(["ts-1"]);
expect(sentPayload.max_budget).toBe(40000);
});
});
describe("KeyInfoView handleKeyUpdate budget_duration", () => {
it("should send a canonical budget_duration through unchanged", async () => {
renderView(true);

View file

@ -631,6 +631,40 @@ describe("KeyEditView", () => {
});
});
it("should keep mcp_toolsets when saving an edit that does not touch the MCP selector", async () => {
const onSubmitMock = vi.fn().mockResolvedValue(undefined);
const keyDataWithToolset = {
...MOCK_KEY_DATA,
object_permission: {
...MOCK_KEY_DATA.object_permission!,
mcp_toolsets: ["ts-1"],
},
};
renderWithProviders(
<KeyEditView
keyData={keyDataWithToolset}
onCancel={() => {}}
onSubmit={onSubmitMock}
accessToken="test-token"
userID="test-user"
userRole="admin"
premiumUser={false}
/>,
);
await waitFor(() => {
expect(screen.getByText("Save Changes")).toBeInTheDocument();
});
await userEvent.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => {
expect(onSubmitMock).toHaveBeenCalled();
});
expect(onSubmitMock.mock.calls[0][0].mcp_servers_and_groups.toolsets).toEqual(["ts-1"]);
});
it("should submit budget_limits: [] when the last budget window is deleted", async () => {
const onSubmitMock = vi.fn().mockResolvedValue(undefined);
const keyDataWithWindow = {

View file

@ -200,6 +200,7 @@ export function KeyEditView({
mcp_servers_and_groups: {
servers: keyData.object_permission?.mcp_servers || [],
accessGroups: keyData.object_permission?.mcp_access_groups || [],
toolsets: keyData.object_permission?.mcp_toolsets || [],
},
mcp_tool_permissions: keyData.object_permission?.mcp_tool_permissions || {},
agents_and_groups: {
@ -233,6 +234,7 @@ export function KeyEditView({
mcp_servers_and_groups: {
servers: keyData.object_permission?.mcp_servers || [],
accessGroups: keyData.object_permission?.mcp_access_groups || [],
toolsets: keyData.object_permission?.mcp_toolsets || [],
},
mcp_tool_permissions: keyData.object_permission?.mcp_tool_permissions || {},
throttle_on_budget_exceeded: keyData.metadata?.throttle_on_budget_exceeded || false,