This commit is contained in:
yuneng-jiang 2026-09-29 14:58:52 -04:00 • committed by GitHub
commit 3a5bb0f22a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 138 additions and 29 deletions

View file

@ -18,8 +18,11 @@ const test = base.extend<{ scopedKey: ScopedKey }>({
team_id: E2E_TEAM_CRUD_ID,
models: [CHAT_MODEL_A],
});
await use({ alias, token: created.token });
await deleteVirtualKey(page.request, created.token);
try {
await use({ alias, token: created.token });
} finally {
await deleteVirtualKey(page.request, created.token);
}
},
});
@ -34,15 +37,44 @@ test.describe("Proxy Admin - Key budget window", () => {
await navigateToPage(page, Page.ApiKeys);
await dismissFeedbackPopup(page);
await openKeyDetail(page, alias);
await page.getByRole("tab", { name: "Settings" }).click();
await page.getByRole("button", { name: "Edit Settings" }).click();
await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("12.5");
// The initial, unfiltered list must already contain this newly created key.
// useKeys keeps that list visible while the filtered search is pending.
await expect(page.getByRole("row").filter({ hasText: alias })).toBeVisible({ timeout: 15_000 });
// Deliver the real search result after editing starts: replacing the list's
// key object must not reset the open form and erase its unsaved budget.
const searchReceived = Promise.withResolvers<void>();
const releaseSearch = Promise.withResolvers<void>();
const isAliasSearch = (url: string) => {
const parsed = new URL(url);
return parsed.pathname.endsWith("/key/list") && parsed.searchParams.get("search") === alias;
};
await page.route("**/key/list?**", async (route) => {
if (!isAliasSearch(route.request().url())) return route.continue();
const response = await route.fetch();
searchReceived.resolve();
await releaseSearch.promise;
await route.fulfill({ response });
});
try {
await openKeyDetail(page, alias);
await page.getByRole("tab", { name: "Settings" }).click();
await page.getByRole("button", { name: "Edit Settings" }).click();
await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("12.5");
await searchReceived.promise;
const refreshed = page.waitForResponse((response) => isAliasSearch(response.url()));
releaseSearch.resolve();
await refreshed;
} finally {
releaseSearch.resolve();
await page.unrouteAll({ behavior: "wait" });
}
await page.getByLabel("Reset Budget", { exact: true }).click();
await page.getByRole("option", { name: "monthly", exact: true }).click();
await page.getByRole("button", { name: "Save Changes" }).click();
const saved = await captureRequestBody(page, { method: "POST", urlIncludes: "/key/update" }, async () => {
await page.getByRole("button", { name: "Save Changes" }).click();
});
expect(saved.max_budget, "a delayed search response must not erase the budget draft").toBe("12.5");
expect(saved.budget_duration).toBe("30d");
await expect
.poll(async () => (await readKeyInfo(page.request, token)).max_budget, {

View file

@ -1004,6 +1004,73 @@ describe("KeyEditView", () => {
});
});
it("should preserve a budget draft when the same key refreshes before saving", async () => {
const onSubmit = vi.fn().mockResolvedValue(undefined);
const props = {
keyData: MOCK_KEY_DATA,
onCancel: vi.fn(),
onSubmit,
accessToken: "test-token",
userID: "test-user",
userRole: "admin",
};
const view = renderWithProviders(<KeyEditView {...props} />);
fireEvent.change(await screen.findByRole("spinbutton", { name: "Max Budget (USD)" }), {
target: { value: "12.5" },
});
view.rerender(
<KeyEditView
{...props}
keyData={{
...MOCK_KEY_DATA,
max_budget: 99,
rpm_limit: 75,
models: ["gpt-4"],
metadata: { note: "new" },
object_permission: { ...MOCK_KEY_DATA.object_permission, vector_stores: ["store-new"] },
}}
/>,
);
await chooseSelectOption(userEvent, screen.getByLabelText("Reset Budget"), "monthly");
await userEvent.click(screen.getByRole("button", { name: /save changes/i }));
const expected = {
max_budget: "12.5",
budget_duration: "30d",
rpm_limit: 75,
models: ["gpt-4"],
metadata: JSON.stringify({ note: "new" }, null, 2),
vector_stores: ["store-new"],
};
await waitFor(() => expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining(expected)));
});
it("should start a fresh budget draft when switching to another key", async () => {
const onSubmit = vi.fn().mockResolvedValue(undefined);
const props = {
keyData: MOCK_KEY_DATA,
onCancel: vi.fn(),
onSubmit,
accessToken: "test-token",
userID: "test-user",
userRole: "admin",
};
const view = renderWithProviders(<KeyEditView {...props} />);
fireEvent.change(await screen.findByRole("spinbutton", { name: "Max Budget (USD)" }), {
target: { value: "12.5" },
});
view.rerender(
<KeyEditView
{...props}
keyData={{ ...MOCK_KEY_DATA, token: "different-key", token_id: "different-key", max_budget: 25 }}
/>,
);
await userEvent.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ max_budget: 25 })));
});
it("should persist a canonical budget_duration value, not a word-form the backend cannot parse", async () => {
const onSubmitMock = vi.fn().mockResolvedValue(undefined);
renderWithProviders(
@ -1545,7 +1612,7 @@ describe("KeyEditView", () => {
});
});
it("should save an explicit project detach while keeping parents locked until the saved key changes", async () => {
it("should save an explicit project detach while keeping parents locked until reopening the editor", async () => {
vi.mocked(getUiSettings).mockResolvedValue({ values: { enable_projects_ui: true } });
const onSubmit = vi.fn().mockResolvedValue(undefined);
const onCancel = vi.fn();
@ -1575,17 +1642,23 @@ describe("KeyEditView", () => {
await userEvent.click(screen.getByRole("button", { name: "Cancel" }));
expect(onCancel).toHaveBeenCalledOnce();
expect(onSubmit).not.toHaveBeenCalled();
// KeyInfoView unmounts the editor on cancel and after a successful save.
view.rerender(<></>);
view.rerender(renderEditor({ ...key }));
await userEvent.click(await screen.findByRole("button", { name: "Detach from project" }));
view.rerender(renderEditor({ ...key, spend: 1 }));
expect(screen.getByRole("button", { name: "Keep project" })).toBeVisible();
await userEvent.click(screen.getByRole("button", { name: /save changes/i }));
const expectedDetach = { project_id: null, organization_id: "org-1", team_id: "group-maple", models: key.models };
await waitFor(() =>
expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining(expectedDetach), expect.any(Array)),
);
expect(screen.getByRole("combobox", { name: "Team ID" })).toBeDisabled();
view.rerender(<></>);
view.rerender(renderEditor({ ...key, project_id: null }));
expect(screen.getByRole("combobox", { name: "Team ID" })).toBeEnabled();
expect(screen.queryByRole("button", { name: "Detach from project" })).not.toBeInTheDocument();
view.rerender(<></>);
view.rerender(renderEditor(key, "Internal User"));
expect(screen.queryByRole("button", { name: "Detach from project" })).not.toBeInTheDocument();
view.rerender(renderEditor(key, "Org Admin"));

View file

@ -88,7 +88,11 @@ interface KeyEditViewProps {
premiumUser?: boolean;
}
export function KeyEditView({
export function KeyEditView(props: KeyEditViewProps) {
return <KeyEditSession key={props.keyData.token} {...props} />;
}
function KeyEditSession({
keyData,
onCancel,
onSubmit,
@ -104,7 +108,9 @@ export function KeyEditView({
const canEditEstimates = userRole != null && isProxyAdminRole(userRole);
const estimateTooltip = estimateTooltips(canEditEstimates);
const form = useZodForm<KeyEditFormValues, KeyEditFormValues>(keyEditFormSchema, {
defaultValues: toKeyEditFormValues(keyData),
// Refresh untouched fields without erasing the user's in-progress edits.
values: toKeyEditFormValues(keyData),
resetOptions: { keepDirtyValues: true },
});
const [promptsList, setPromptsList] = useState<string[]>([]);
const [tagsList, setTagsList] = useState<Record<string, Tag>>({});
@ -187,21 +193,17 @@ export function KeyEditView({
// Sync disabled callbacks with form when component mounts
useEffect(() => {
form.setValue("disabled_callbacks", disabledCallbacks);
form.setValue("disabled_callbacks", disabledCallbacks, { shouldDirty: true });
}, [form, disabledCallbacks]);
useEffect(() => {
form.reset(toKeyEditFormValues(keyData));
}, [keyData, form]);
// Sync auto-rotation state with form values
useEffect(() => {
form.setValue("auto_rotate", autoRotationEnabled);
form.setValue("auto_rotate", autoRotationEnabled, { shouldDirty: true });
}, [autoRotationEnabled, form]);
useEffect(() => {
if (rotationInterval) {
form.setValue("rotation_interval", rotationInterval);
form.setValue("rotation_interval", rotationInterval, { shouldDirty: true });
}
}, [rotationInterval, form]);
@ -327,7 +329,7 @@ export function KeyEditView({
const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | null) => {
setField(orgId);
setSelectedOrganizationId(orgId);
form.setValue("team_id", null);
form.setValue("team_id", null, { shouldDirty: true });
};
const handleTeamChange = (setField: (value: string | null) => void, teamId: string | null) => {
@ -335,16 +337,16 @@ export function KeyEditView({
const selectedTeam = teams?.find((t) => t.team_id === teamId) || null;
if (selectedTeam?.organization_id) {
setSelectedOrganizationId(selectedTeam.organization_id);
form.setValue("organization_id", selectedTeam.organization_id);
form.setValue("organization_id", selectedTeam.organization_id, { shouldDirty: true });
} else if (!teamId) {
setSelectedOrganizationId(null);
form.setValue("organization_id", null);
form.setValue("organization_id", null, { shouldDirty: true });
}
};
const handleDisabledCallbacksChange = (internalValues: string[]) => {
setDisabledCallbacks(mapInternalToDisplayNames(internalValues));
form.setValue("disabled_callbacks", internalValues);
form.setValue("disabled_callbacks", internalValues, { shouldDirty: true });
};
const modelOptions = [
@ -409,14 +411,14 @@ export function KeyEditView({
onChange={(value) => {
switch (value) {
case "default":
form.setValue("allowed_routes", "");
form.setValue("allowed_routes", "", { shouldDirty: true });
break;
case "llm_api":
form.setValue("allowed_routes", "llm_api_routes");
form.setValue("allowed_routes", "llm_api_routes", { shouldDirty: true });
break;
case "management":
form.setValue("allowed_routes", "management_routes");
form.setValue("models", []);
form.setValue("allowed_routes", "management_routes", { shouldDirty: true });
form.setValue("models", [], { shouldDirty: true });
break;
}
}}
@ -766,7 +768,7 @@ export function KeyEditView({
selectedAccessGroups={mcpSelection?.accessGroups || []}
selectedToolsets={mcpSelection?.toolsets || []}
toolPermissions={(mcpToolPermissions as Record<string, string[]> | undefined) || {}}
onChange={(toolPerms) => form.setValue("mcp_tool_permissions", toolPerms)}
onChange={(toolPerms) => form.setValue("mcp_tool_permissions", toolPerms, { shouldDirty: true })}
/>
</div>
@ -828,7 +830,9 @@ export function KeyEditView({
canDetach={canDetachProject}
pending={detachProject}
disabled={isKeySaving}
onToggle={() => form.setValue("project_id", detachProject ? keyData.project_id : null)}
onToggle={() =>
form.setValue("project_id", detachProject ? keyData.project_id : null, { shouldDirty: true })
}
/>
)}