fix(usage): scope key search to the resolved user and reset stale remote results

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
(cherry picked from commit 3e74d6d613)

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -3171,13 +3171,16 @@ async def search_user_daily_activity_keys(
     try:
         entity_id: Final = _resolve_user_daily_activity_entity_id(user_api_key_dict, user_id)
 
-        where: Final[KeyActivitySearchWhere] = {
-            "OR": (
-                {"token": search},
-                {"key_alias": {"contains": search, "mode": "insensitive"}},
-                {"user_id": {"contains": search, "mode": "insensitive"}},
-            )
-        }
+        search_or: Final = (
+            {"token": search},  # mutable-ok: prisma serializes where clauses, keep plain dicts
+            {"key_alias": {"contains": search, "mode": "insensitive"}},  # mutable-ok: prisma where clause leaf
+            {"user_id": {"contains": search, "mode": "insensitive"}},  # mutable-ok: prisma where clause leaf
+        )
+        where: Final[KeyActivitySearchWhere] = (
+            {"OR": search_or}  # mutable-ok: prisma where clause root
+            if entity_id is None
+            else {"user_id": entity_id, "OR": search_or}  # mutable-ok: prisma where clause root
+        )
         matched_keys: Final = await VerificationTokenRepository(prisma_client).table.find_many(
             where=where,
             take=USAGE_TOP_API_KEYS_LIMIT,

litellm/types/proxy/management_endpoints/internal_user_endpoints.py

@@ -2,7 +2,7 @@ from collections.abc import Mapping, Sequence
 from typing import Any, Final, Literal
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator
-from typing_extensions import ReadOnly, TypedDict
+from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 from litellm.proxy._types import (
     LiteLLM_UserTableWithKeyCount,
@@ -32,6 +32,7 @@ class KeyActivitySearchWhere(TypedDict):
     """Prisma filter behind `/user/daily/activity/aggregated/search`: exact token hash, or key alias
     or user id containing the term, case-insensitive."""
 
+    user_id: NotRequired[ReadOnly[str]]
     OR: ReadOnly[
         tuple[Mapping[Literal["token"], str] | Mapping[Literal["key_alias", "user_id"], InsensitiveContains], ...]
     ]

tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py

@@ -2709,6 +2709,7 @@ async def test_search_user_daily_activity_keys_passes_matched_tokens_to_aggregat
         {"key_alias": {"contains": "gamma", "mode": "insensitive"}},
         {"user_id": {"contains": "gamma", "mode": "insensitive"}},
     )
+    assert "user_id" not in find_many_kwargs["where"]
 
     mock_get_daily_agg.assert_called_once_with(
         prisma_client=mock_prisma_client,
@@ -2806,6 +2807,8 @@ async def test_search_user_daily_activity_keys_non_admin_scoped_to_caller(monkey
 
     assert result is mock_response
     assert mock_get_daily_agg.call_args.kwargs["entity_id"] == "user-1"
+    find_many_kwargs = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs
+    assert find_many_kwargs["where"]["user_id"] == "user-1"
 
     with pytest.raises(HTTPException) as exc_info:
         await search_user_daily_activity_keys(

ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx

@@ -107,6 +107,28 @@ describe("KeyActivityPanel", () => {
     expect(searchKeys).not.toHaveBeenCalled();
   });
 
+  it("drops stale server results as soon as the search callback is rebuilt", async () => {
+    const searchKeysA = vi
+      .fn<(query: string) => Promise<Record<string, ModelActivityData>>>()
+      .mockResolvedValue({ "hash-gamma": activity("gamma-low-key", "gamma@example.com", "user-gamma") });
+    const searchKeysB = vi
+      .fn<(query: string) => Promise<Record<string, ModelActivityData>>>()
+      .mockReturnValue(new Promise(() => {}));
+    const { rerender } = render(
+      <KeyActivityPanel keyMetrics={keyMetrics} apiKeyTruncation={{ limit: 2, total: 3 }} searchKeys={searchKeysA} />,
+    );
+
+    fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "gamma" } });
+    expect(await screen.findByText("hash-gamma")).toBeInTheDocument();
+
+    rerender(
+      <KeyActivityPanel keyMetrics={keyMetrics} apiKeyTruncation={{ limit: 2, total: 3 }} searchKeys={searchKeysB} />,
+    );
+
+    expect(screen.getByRole("status")).toHaveTextContent("Searching all keys");
+    expect(screen.queryByText("hash-gamma")).not.toBeInTheDocument();
+  });
+
   it("reports a failed server search but keeps the local matches", async () => {
     const searchKeys = vi
       .fn<(query: string) => Promise<Record<string, ModelActivityData>>>()

ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx

@@ -12,14 +12,16 @@ interface KeyActivityPanelProps {
   keyMetrics: Record<string, ModelActivityData>;
   hidePromptCachingMetrics?: boolean;
   apiKeyTruncation?: ApiKeyTruncation;
-  searchKeys?: (query: string) => Promise<Record<string, ModelActivityData>>;
+  searchKeys?: SearchKeys;
 }
 
+type SearchKeys = (query: string) => Promise<Record<string, ModelActivityData>>;
+
 type RemoteSearch =
   | { status: "idle" }
-  | { status: "loading"; query: string }
-  | { status: "done"; query: string; keys: Record<string, ModelActivityData> }
-  | { status: "error"; query: string };
+  | { status: "loading"; query: string; searchKeys: SearchKeys }
+  | { status: "done"; query: string; searchKeys: SearchKeys; keys: Record<string, ModelActivityData> }
+  | { status: "error"; query: string; searchKeys: SearchKeys };
 
 const REMOTE_SEARCH_DEBOUNCE_MS = 300;
 
@@ -39,13 +41,13 @@ const KeyActivityPanel: React.FC<KeyActivityPanelProps> = ({
     if (!remoteEnabled) return;
     let cancelled = false;
     const timer = setTimeout(() => {
-      setRemote({ status: "loading", query: trimmedQuery });
+      setRemote({ status: "loading", query: trimmedQuery, searchKeys });
       searchKeys(trimmedQuery)
         .then((keys) => {
-          if (!cancelled) setRemote({ status: "done", query: trimmedQuery, keys });
+          if (!cancelled) setRemote({ status: "done", query: trimmedQuery, searchKeys, keys });
         })
         .catch(() => {
-          if (!cancelled) setRemote({ status: "error", query: trimmedQuery });
+          if (!cancelled) setRemote({ status: "error", query: trimmedQuery, searchKeys });
         });
     }, REMOTE_SEARCH_DEBOUNCE_MS);
     return () => {
@@ -54,8 +56,10 @@ const KeyActivityPanel: React.FC<KeyActivityPanelProps> = ({
     };
   }, [remoteEnabled, trimmedQuery, searchKeys]);
 
-  const remoteCurrent = remoteEnabled && "query" in remote && remote.query === trimmedQuery;
-  const remoteLoading = remoteEnabled && !remoteCurrent;
+  const remoteMatchesSearch =
+    "searchKeys" in remote && remote.searchKeys === searchKeys && remote.query === trimmedQuery;
+  const remoteCurrent = remoteEnabled && remoteMatchesSearch;
+  const remoteLoading = remoteEnabled && (remote.status === "loading" || !remoteCurrent);
   const remoteFailed = remoteCurrent && remote.status === "error";
 
   const extraRemoteKeys = useMemo(() => {

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

@@ -2562,9 +2562,6 @@ export const userDailyActivityKeySearchCall = async (
   endTime: Date,
   ...options: [search: string, userId?: string | null]
 ) => {
-  /**
-   * Search keys beyond the top-spend subset the aggregated endpoint loads
-   */
   const [search, userId = null] = options;
   try {
     const formatDate = (date: Date) => {