diff --git a/litellm/proxy/memory/management.py b/litellm/proxy/memory/management.py index 483f4f8c225..d0974f1f543 100644 --- a/litellm/proxy/memory/management.py +++ b/litellm/proxy/memory/management.py @@ -1,6 +1,7 @@ from typing import Annotated, Final from fastapi import APIRouter, Depends, HTTPException, Query, Response +from pydantic import AwareDatetime from litellm.proxy._types import UI_TEAM_ID, LitellmUserRoles, UserAPIKeyAuth, user_api_key_has_admin_view from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -259,11 +260,18 @@ async def list_entries( offset: int = Query(0, ge=0), key_id: str | None = Query(None, pattern=r"^[a-f0-9]{64}$"), auth: UserAPIKeyAuth = _AUTH, + before_updated_at: Annotated[AwareDatetime | None, Query()] = None, + before_memory_id: Annotated[str | None, Query(min_length=1, max_length=128)] = None, ) -> tuple[MemoryEntry, ...]: + if (before_updated_at is None) != (before_memory_id is None): + raise HTTPException(status_code=422, detail="Provide both memory cursor fields") prisma: Final = memory_primary_client(require_memory_prisma()) access: Final = await access_for_key(auth, key_id) return await MemoryStore(prisma, access).search( - MemorySearch(query=query, limit=limit, offset=offset), require_active=False, recent_first=True + MemorySearch(query=query, limit=limit, offset=offset), + require_active=False, + recent_first=True, + before=(before_updated_at, before_memory_id) if before_updated_at and before_memory_id else None, ) diff --git a/litellm/proxy/memory/policy.py b/litellm/proxy/memory/policy.py index cebec0ac27c..12d4a7f3895 100644 --- a/litellm/proxy/memory/policy.py +++ b/litellm/proxy/memory/policy.py @@ -169,7 +169,7 @@ class MemoryAccess: return MemoryStatus( active=self.active, activation=self.policy.activation if self.policy else "disabled", - scope=self.policy.scope if self.policy else None, + scope=self.policy.scope if self.policy and self.namespace else None, opted_in=self.opted_in, policy_id=self.policy.policy_id if self.policy else None, ) diff --git a/litellm/proxy/memory/store.py b/litellm/proxy/memory/store.py index 6c10c429af9..2295e591c1b 100644 --- a/litellm/proxy/memory/store.py +++ b/litellm/proxy/memory/store.py @@ -1,5 +1,6 @@ import asyncio import json +from datetime import datetime from types import SimpleNamespace from typing import TYPE_CHECKING, Final @@ -68,7 +69,12 @@ class MemoryStore: return current.namespace async def search( - self, search: MemorySearch, *, require_active: bool = True, recent_first: bool = False + self, + search: MemorySearch, + *, + require_active: bool = True, + recent_first: bool = False, + before: tuple[datetime, str] | None = None, ) -> tuple[MemoryEntry, ...]: entries: Final = await self.entries(require_active=require_active) ranked: Final = await asyncio.to_thread(fuzzy_memories, search.query, entries) @@ -78,7 +84,15 @@ class MemoryStore: if recent_first else tuple(entry for entry, _, _ in ranked) ) - return ordered[search.offset : search.offset + search.limit] + page: Final = tuple( + entry + for entry in ordered + if before is None + or entry.updated_at < before[0] + or entry.updated_at == before[0] + and entry.memory_id > before[1] + ) + return page[search.offset : search.offset + search.limit] async def entries(self, *, require_active: bool = True) -> tuple[MemoryEntry, ...]: namespace: Final = await self.authorize_namespace(require_active=require_active) diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_management.py b/tests/test_litellm/proxy/memory/test_memory_v2_management.py index 3b4fc918016..73b908dd92d 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_management.py +++ b/tests/test_litellm/proxy/memory/test_memory_v2_management.py @@ -354,3 +354,55 @@ async def test_dashboard_search_keeps_recency_before_pagination_while_agent_sear "older", "newer", ] + + +@pytest.mark.asyncio +async def test_dashboard_cursor_survives_deletion_of_earlier_entries(database: MagicMock) -> None: + database.db.litellm_memorypolicy.find_many.return_value = [policy()] + now = datetime(2026, 9, 12, tzinfo=timezone.utc) + rows = [ + SimpleNamespace( + memory_id=identifier, + key=f"memory-v2:namespace:{identifier}", + value=f"Memory {identifier}", + metadata={"title": identifier}, + updated_at=now, + created_at=now, + created_by="owner", + ) + for identifier in ("a", "b", "c") + ] + database.db.litellm_memorytable.find_many.return_value = rows + first = await management.list_entries("", 1, 0, None, auth()) + assert first[0].memory_id == "a" + database.db.litellm_memorytable.find_many.return_value = rows[1:] + second = await management.list_entries("", 1, 0, None, auth(), first[0].updated_at, first[0].memory_id) + assert second[0].memory_id == "b" + third = await management.list_entries("", 1, 0, None, auth(), second[0].updated_at, second[0].memory_id) + assert third[0].memory_id == "c" + + +@pytest.mark.parametrize( + "params", + [ + {"before_memory_id": "a"}, + {"before_updated_at": "2026-09-12T00:00:00Z"}, + {"before_memory_id": "a", "before_updated_at": "2026-09-12T00:00:00"}, + ], +) +def test_dashboard_cursor_rejects_partial_or_naive_dates(database: MagicMock, params: dict[str, str]) -> None: + app = FastAPI() + app.include_router(management.router) + app.dependency_overrides[user_api_key_auth] = auth + with TestClient(app) as client: + assert client.get("/v2/memory/entries", params=params).status_code == 422 + database.db.litellm_memorytable.find_many.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("activation", ["automatic", "opt_in"]) +async def test_status_does_not_offer_an_unresolvable_namespace(database: MagicMock, activation: str) -> None: + database.db.litellm_verificationtoken.find_unique.return_value["user_id"] = None + database.db.litellm_memorypolicy.find_many.return_value = [policy(activation=activation, scope="user")] + status = await management.get_status("a" * 64, auth()) + assert not status.active and status.scope is None diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts index 7e7089e685f..8430befb15a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts @@ -33,6 +33,8 @@ export interface DeletedKeysResponse { } export interface KeyListCallOptions { + includeTeamKeys?: boolean; + includeCreatedByKeys?: boolean; organizationID?: string | null; teamID?: string | null; projectID?: string | null; @@ -71,8 +73,8 @@ const keyListCall = async (accessToken: string, page: number, pageSize: number, expand: options.expand, status: options.status, return_full_object: "true", - include_team_keys: "true", - include_created_by_keys: "true", + include_team_keys: options.includeTeamKeys ?? true, + include_created_by_keys: options.includeCreatedByKeys ?? true, // Opt into substring matching so the admin key-list search box keeps // matching partial user_id/key_alias. /key/list is exact by default. substring_matching: "true", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/AutomaticMemoryEntries.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/AutomaticMemoryEntries.tsx index c99788ee782..e17f8d7c80a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/AutomaticMemoryEntries.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/AutomaticMemoryEntries.tsx @@ -22,13 +22,17 @@ import { MemoryKeyPicker } from "./MemoryTargetPicker"; type Entry = components["schemas"]["MemoryEntry"]; type Capture = components["schemas"]["MemoryCapture"]; +type Cursor = { before_updated_at: string; before_memory_id: string } | null; type Status = components["schemas"]["MemoryStatus"]; function memoryDescription(status?: Status) { if (!status) return "What your assistants remember across conversations."; - if (status.activation === "automatic") return "Memory is on for this key. Your administrator manages this setting."; if (status.activation === "disabled" || !status.scope) return "Memory is off. Your administrator can make it available for this key."; + if (status.activation === "automatic") { + if (status.active) return "Memory is on for this key. Your administrator manages this setting."; + return "Memory is off. Your administrator can make it available for this key."; + } if (status.active) return "Your assistants can save and recall memories. You can turn this off at any time."; return "Your assistants won't save or recall memories. Turn it on whenever you're ready."; } @@ -63,7 +67,14 @@ type DashboardProps = Readonly<{ userId: string; readOnly: boolean; proxyAdmin: export function AutomaticMemoryEntries({ userId, readOnly, proxyAdmin }: DashboardProps) { const [selection, setSelection] = useState(); - const keys = useKeys(1, 1, { userID: proxyAdmin ? undefined : userId, sortBy: "created_at", sortOrder: "desc" }); + const keyOptions = { + userID: proxyAdmin ? undefined : userId, + sortBy: "created_at", + sortOrder: "desc", + includeTeamKeys: proxyAdmin, + includeCreatedByKeys: proxyAdmin, + }; + const keys = useKeys(1, 1, keyOptions); const keyId = selection ?? keys.data?.keys[0]?.token ?? ""; return ( @@ -115,16 +126,20 @@ function MemoryDashboard({ const entriesOptions = { queryKey: ["memoryEntries", userId, keyId, query], enabled: !!keyId && !!status.data?.scope, - initialPageParam: 0, - queryFn: async ({ signal, pageParam }: { signal: AbortSignal; pageParam: number }) => + initialPageParam: null as Cursor, + queryFn: async ({ signal, pageParam }: { signal: AbortSignal; pageParam: Cursor }) => ( await fetchClient.GET("/v2/memory/entries", { - params: { query: { key_id: keyId, query, offset: pageParam, limit: 20 } }, + params: { query: { key_id: keyId, query, limit: 20, ...pageParam } }, signal, }) ).data ?? [], - getNextPageParam: (lastPage: Entry[], _pages: Entry[][], offset: number) => - lastPage.length === 20 ? offset + 20 : undefined, + getNextPageParam: (lastPage: Entry[]) => { + const last = lastPage.at(-1); + return lastPage.length === 20 && last + ? { before_updated_at: last.updated_at, before_memory_id: last.memory_id } + : undefined; + }, }; const entries = useInfiniteQuery(entriesOptions); const save = useMutation({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTargetPicker.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTargetPicker.tsx index 152ab00c20d..4a9c042a8d7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTargetPicker.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTargetPicker.tsx @@ -68,7 +68,8 @@ export function MemoryKeyPicker({ userId, }: PickerProps & Readonly<{ inputId?: string; userId?: string }>) { const [search, setSearch] = useState(""); - const query = useInfiniteKeys(25, { search, userID: userId }); + const keyOptions = { search, userID: userId, includeTeamKeys: !userId, includeCreatedByKeys: !userId }; + const query = useInfiniteKeys(25, keyOptions); return ( (); const calls: { path: string; method: string; body: unknown; keyId: string | null }[] = []; let enabled = false; let activation = "opt_in"; +let available = true; let paginated = false; let failPreference = false; const entry = { @@ -41,6 +42,7 @@ beforeEach(async () => { calls.length = 0; enabled = false; activation = "opt_in"; + available = true; paginated = false; failPreference = false; vi.clearAllMocks(); @@ -49,7 +51,12 @@ beforeEach(async () => { input instanceof Request ? input : new Request(new URL(String(input), window.location.origin), init); const path = new URL(request.url).pathname; const text = request.method === "GET" ? "" : await request.text(); - const keyId = new URL(request.url).searchParams.get("key_id"); + const url = new URL(request.url); + const keyId = url.searchParams.get("key_id"); + if (url.pathname === "/key/list" && url.searchParams.get("user_id")) { + expect(url.searchParams.get("include_team_keys")).toBe("false"); + expect(url.searchParams.get("include_created_by_keys")).toBe("false"); + } const call = { path, method: request.method, body: text ? JSON.parse(text) : undefined, keyId }; calls.push(call); if (path === "/v2/memory/preference" && request.method === "PUT" && failPreference) { @@ -66,13 +73,18 @@ beforeEach(async () => { if (path === "/v2/memory/policies") return []; if (path === "/v1/memory") return { memories: [], total: 0 }; if (path === "/v2/memory/status") - return { active: activation === "automatic" || enabled, opted_in: enabled, activation, scope: "key" }; + return { + active: available && (activation === "automatic" || enabled), + opted_in: enabled, + activation, + scope: available ? "key" : null, + }; if (path === "/v2/memory/entries") { if (request.method === "POST") return entry; if (keyId === "b".repeat(64)) return [{ ...entry, memory_id: "other", title: "Other key memory", content: "Another project" }]; if (paginated) { - const offset = Number(new URL(request.url).searchParams.get("offset")); + const offset = new URL(request.url).searchParams.get("before_memory_id") === "entry-19" ? 20 : 0; return Array.from({ length: offset ? 1 : 20 }, (_, i) => ({ ...entry, memory_id: `entry-${offset + i}`, @@ -211,6 +223,17 @@ describe("Memory dashboard", () => { expect(screen.getByText(/Your administrator manages this setting/)).toBeVisible(); }); + it("shows memory as unavailable when an automatic policy cannot resolve a sharing scope", async () => { + session("internal_user"); + activation = "automatic"; + available = false; + renderWithProviders(); + const toggle = await screen.findByRole("switch", { name: "Memory" }); + expect(toggle).not.toBeChecked(); + expect(toggle).toHaveAttribute("aria-disabled", "true"); + expect(screen.getByText("Memory is off. Your administrator can make it available for this key.")).toBeVisible(); + }); + it("keeps the actual state off when saving a preference fails", async () => { session("internal_user"); failPreference = true; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/page.tsx index 120d2e68fcc..9fa4796dba3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/page.tsx @@ -30,7 +30,7 @@ export default function Memory() { return (
- {userId && } + {userId && } {canManage && userId && ( }> diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d7e58550bfa..04571228e30 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -67301,6 +67301,8 @@ export interface operations { limit?: number; offset?: number; key_id?: string | null; + before_updated_at?: string | null; + before_memory_id?: string | null; }; header?: never; path?: never;