From 104d81cabadd3cf835affb979c11c2efcca42121 Mon Sep 17 00:00:00 2001
From: willcai1984 <willcai1984@gmail.com>
Date: Fri, 7 Aug 2026 18:50:36 +0800
Subject: [PATCH] fix(ui): invalidate keys list after key update so table
 reflects changes

handleKeyUpdate in key_info_view.tsx updated local component state and called
onKeyDataUpdate after a successful /key/update, but never invalidated the
React Query keys list -- unlike handleDelete in the same file, which already
calls queryClient.invalidateQueries({ queryKey: keyKeys.lists() }).

As a result, editing a key (e.g. daily budget / budget reset) and returning
to the keys table showed stale cached values: the change looked like it "didn't
save" even though /key/update had persisted it and refreshed the runtime key
cache on the backend.

Fix: invalidate keyKeys.lists() after a successful update, mirroring the
delete path. Adds a regression test that fails before the fix (spy never
called with { queryKey: ['keys','list'] }) and passes after.
---
 .../KeyInfoView.handleKeyUpdate.test.tsx      | 32 +++++++++++++++++--
 .../components/templates/key_info_view.tsx    |  3 +-
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx
index 16b1be0d74a..30bb8c22f3e 100644
--- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx
+++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx
@@ -2,11 +2,12 @@ import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 
 // ---- Hoisted shared mocks (safe to use inside vi.mock factories) ----
-const { keyUpdateCallMock, keyDeleteCallMock, mockUseAuthorized } = vi.hoisted(() => {
+const { keyUpdateCallMock, keyDeleteCallMock, mockUseAuthorized, invalidateQueriesMock } = vi.hoisted(() => {
   return {
     keyUpdateCallMock: vi.fn().mockResolvedValue({}),
     keyDeleteCallMock: vi.fn().mockResolvedValue({}),
     mockUseAuthorized: vi.fn(),
+    invalidateQueriesMock: vi.fn(),
   };
 });
 
@@ -270,12 +271,13 @@ vi.mock("@/app/(dashboard)/hooks/keys/useSetKeyBlockedState", () => ({
   }),
 }));
 
-// useQueryClient also needs a provider; the delete-path invalidation is covered in key_info_view.test.tsx
+// useQueryClient also needs a provider; return a hoisted spy so we can assert the
+// update path invalidates the keys list (mirrors the delete-path test in key_info_view.test.tsx)
 vi.mock("@tanstack/react-query", async (importOriginal) => {
   const actual = await importOriginal<typeof import("@tanstack/react-query")>();
   return {
     ...actual,
-    useQueryClient: () => ({ invalidateQueries: vi.fn() }),
+    useQueryClient: () => ({ invalidateQueries: invalidateQueriesMock }),
   };
 });
 
@@ -301,6 +303,7 @@ vi.mock("./key_edit_view", async () => {
 
 // ---- SUT import AFTER mocks ----
 import KeyInfoView from "./key_info_view";
+import { keyKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
 
 // ---- Test data helpers ----
 const baseKeyData = {
@@ -593,3 +596,26 @@ describe("KeyInfoView handleKeyUpdate empty strings", () => {
     });
   });
 });
+
+describe("KeyInfoView handleKeyUpdate keys-list invalidation", () => {
+  it("invalidates the keys list query after a successful update so the table reflects the change", 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: 60,
+      budget_duration: "24h",
+    };
+
+    fireEvent.click(screen.getByText("Mock Submit"));
+
+    await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled());
+    // The keys list must be invalidated, otherwise returning to the table shows
+    // stale (cached) values — the original bug behind "budget save looks invalid".
+    await waitFor(() => {
+      expect(invalidateQueriesMock).toHaveBeenCalledWith({ queryKey: keyKeys.lists() });
+    });
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx
index f4803727875..21c99f57777 100644
--- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx
+++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx
@@ -314,7 +314,8 @@ export default function KeyInfoView({
       }
       NotificationManager.success("Key updated successfully");
       setIsEditing(false);
-      // Refresh key data here if needed
+      // Invalidate the keys list so the table reflects the update (mirrors handleDelete)
+      await queryClient.invalidateQueries({ queryKey: keyKeys.lists() });
     } catch (error) {
       NotificationManager.fromBackend(parseErrorMessage(error));
       console.error("Error updating key:", error);