diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index f23956d5c5b..36e10eddce7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1224,9 +1224,13 @@ model LiteLLM_ClaudeCodePluginTable { @@map("LiteLLM_ClaudeCodePluginTable") } -// User/team-scoped memory store. Keyed by (key, user_id, team_id). +// User/team-scoped memory store with a GLOBAL unique key. // `value` is a string (typically markdown/text meant for LLM context); // `metadata` is an optional JSON envelope for structured tags without schema changes. +// Note: `key` is globally unique across all users/teams — callers should +// namespace their keys (e.g. `user:123:notes`) if they need per-user isolation. +// `user_id` / `team_id` stamp ownership for visibility filtering, but do NOT +// participate in the unique constraint. model LiteLLM_MemoryTable { memory_id String @id @default(uuid()) key String @unique diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index bdc73046c8c..7e6cfc3a645 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -86,6 +86,67 @@ def _require_prisma(): return prisma_client +def _is_unique_violation(exc: Exception) -> bool: + """Best-effort detection of a Prisma unique-constraint violation.""" + msg = str(exc) + return "Unique" in msg or "unique" in msg or "UniqueViolation" in msg + + +def _resolve_scope( + user_api_key_dict: UserAPIKeyAuth, + requested_user_id: Optional[str], + requested_team_id: Optional[str], +) -> tuple[Optional[str], Optional[str]]: + """ + Resolve the (user_id, team_id) to stamp on a new row. + + - PROXY_ADMIN: may override either dimension via the request body. + - Everyone else: the requested values must match their own (or be omitted). + + Also rejects identity-less creation: a row with both user_id and team_id + NULL is invisible to every non-admin caller (the visibility filter would + never match it), so we refuse to create orphan rows unless the caller is + a PROXY_ADMIN who is explicitly stamping a global/shared row. + """ + if _is_admin(user_api_key_dict): + user_id = ( + requested_user_id + if requested_user_id is not None + else user_api_key_dict.user_id + ) + team_id = ( + requested_team_id + if requested_team_id is not None + else user_api_key_dict.team_id + ) + return user_id, team_id + + if requested_user_id is not None and requested_user_id != user_api_key_dict.user_id: + raise HTTPException( + status_code=403, + detail="Only proxy admins may set user_id to a different user.", + ) + if requested_team_id is not None and requested_team_id != user_api_key_dict.team_id: + raise HTTPException( + status_code=403, + detail="Only proxy admins may set team_id to a different team.", + ) + user_id = user_api_key_dict.user_id + team_id = user_api_key_dict.team_id + if not user_id and not team_id: + # Orphan row: no user_id and no team_id means no non-admin can ever + # see it again via the visibility filter. Reject up front. + raise HTTPException( + status_code=400, + detail=( + "Cannot create a memory entry without a user_id or team_id. " + "Authenticate with a key that has a user_id or team_id, or call " + "as a proxy admin." + ), + ) + return user_id, team_id + + @router.post( "/v1/memory", tags=["memory management"], @@ -98,28 +159,7 @@ async def create_memory( ): """Create a new memory entry for the caller (or, for admins, any scope).""" prisma_client = _require_prisma() - - if _is_admin(user_api_key_dict): - user_id = ( - body.user_id if body.user_id is not None else user_api_key_dict.user_id - ) - team_id = ( - body.team_id if body.team_id is not None else user_api_key_dict.team_id - ) - else: - # Non-admins cannot set a scope other than their own. - if body.user_id is not None and body.user_id != user_api_key_dict.user_id: - raise HTTPException( - status_code=403, - detail="Only proxy admins may set user_id to a different user.", - ) - if body.team_id is not None and body.team_id != user_api_key_dict.team_id: - raise HTTPException( - status_code=403, - detail="Only proxy admins may set team_id to a different team.", - ) - user_id = user_api_key_dict.user_id - team_id = user_api_key_dict.team_id + user_id, team_id = _resolve_scope(user_api_key_dict, body.user_id, body.team_id) # Prisma's Python client rejects `metadata=None` on a `Json?` field — # the field must be omitted entirely to store SQL NULL. Build the data @@ -139,8 +179,7 @@ async def create_memory( row = await prisma_client.db.litellm_memorytable.create(data=create_data) except Exception as e: # Key is globally unique. Any duplicate → 409. - msg = str(e) - if "Unique" in msg or "unique" in msg or "UniqueViolation" in msg: + if _is_unique_violation(e): raise HTTPException( status_code=409, detail=f"Memory with key '{body.key}' already exists.", @@ -265,16 +304,17 @@ async def upsert_memory( ) data["updated_by"] = user_api_key_dict.user_id - try: - existing = None + async def _find_existing() -> Any: + """Return the caller-visible row for `key`, or None.""" try: - existing = await _find_memory_for_caller( - prisma_client, key, user_api_key_dict - ) + return await _find_memory_for_caller(prisma_client, key, user_api_key_dict) except HTTPException as e: - if e.status_code != 404: - raise + if e.status_code == 404: + return None + raise + try: + existing = await _find_existing() if existing is not None: row = await prisma_client.db.litellm_memorytable.update( where={"memory_id": existing.memory_id}, @@ -286,18 +326,46 @@ async def upsert_memory( status_code=400, detail="Cannot create a new memory via PUT without a 'value'.", ) + # PUT-create must honor admin scope override, matching POST semantics. + user_id, team_id = _resolve_scope( + user_api_key_dict, body.user_id, body.team_id + ) # Omit `metadata` when None — Prisma rejects None on Json? fields. create_data: dict = { "key": key, "value": body.value, - "user_id": user_api_key_dict.user_id, - "team_id": user_api_key_dict.team_id, + "user_id": user_id, + "team_id": team_id, "created_by": user_api_key_dict.user_id, "updated_by": user_api_key_dict.user_id, } if body.metadata is not None: create_data["metadata"] = body.metadata - row = await prisma_client.db.litellm_memorytable.create(data=create_data) + try: + row = await prisma_client.db.litellm_memorytable.create( + data=create_data + ) + except Exception as e: + # Race: a concurrent PUT/POST created the row after our check. + # Re-read and fall back to an update so the PUT stays idempotent + # instead of surfacing a 500 on a unique-violation. + if not _is_unique_violation(e): + raise + existing_after_race = await _find_existing() + if existing_after_race is None: + # Row exists globally but isn't visible to this caller + # (owned by someone else). Treat as conflict. + raise HTTPException( + status_code=409, + detail=( + f"Memory with key '{key}' already exists in a " + "different scope." + ), + ) + row = await prisma_client.db.litellm_memorytable.update( + where={"memory_id": existing_after_race.memory_id}, + data=data, + ) except HTTPException: raise except Exception as e: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index f23956d5c5b..36e10eddce7 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1224,9 +1224,13 @@ model LiteLLM_ClaudeCodePluginTable { @@map("LiteLLM_ClaudeCodePluginTable") } -// User/team-scoped memory store. Keyed by (key, user_id, team_id). +// User/team-scoped memory store with a GLOBAL unique key. // `value` is a string (typically markdown/text meant for LLM context); // `metadata` is an optional JSON envelope for structured tags without schema changes. +// Note: `key` is globally unique across all users/teams — callers should +// namespace their keys (e.g. `user:123:notes`) if they need per-user isolation. +// `user_id` / `team_id` stamp ownership for visibility filtering, but do NOT +// participate in the unique constraint. model LiteLLM_MemoryTable { memory_id String @id @default(uuid()) key String @unique diff --git a/litellm/types/memory_management.py b/litellm/types/memory_management.py index f400f71313c..bde54933699 100644 --- a/litellm/types/memory_management.py +++ b/litellm/types/memory_management.py @@ -43,6 +43,11 @@ class MemoryCreateRequest(BaseModel): class MemoryUpdateRequest(BaseModel): value: Optional[str] = None metadata: Optional[Any] = None + # Only honored on create (when the row doesn't yet exist) and only for + # PROXY_ADMIN callers — mirrors MemoryCreateRequest so admins can bootstrap + # rows scoped to another user/team via PUT, not just POST. + user_id: Optional[str] = None + team_id: Optional[str] = None class MemoryListResponse(BaseModel): diff --git a/schema.prisma b/schema.prisma index f23956d5c5b..36e10eddce7 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1224,9 +1224,13 @@ model LiteLLM_ClaudeCodePluginTable { @@map("LiteLLM_ClaudeCodePluginTable") } -// User/team-scoped memory store. Keyed by (key, user_id, team_id). +// User/team-scoped memory store with a GLOBAL unique key. // `value` is a string (typically markdown/text meant for LLM context); // `metadata` is an optional JSON envelope for structured tags without schema changes. +// Note: `key` is globally unique across all users/teams — callers should +// namespace their keys (e.g. `user:123:notes`) if they need per-user isolation. +// `user_id` / `team_id` stamp ownership for visibility filtering, but do NOT +// participate in the unique constraint. model LiteLLM_MemoryTable { memory_id String @id @default(uuid()) key String @unique diff --git a/tests/test_litellm/proxy/memory/test_memory_endpoints.py b/tests/test_litellm/proxy/memory/test_memory_endpoints.py index 9733b5d9ec8..90863e5e492 100644 --- a/tests/test_litellm/proxy/memory/test_memory_endpoints.py +++ b/tests/test_litellm/proxy/memory/test_memory_endpoints.py @@ -203,6 +203,76 @@ class TestMemoryEndpoints: ) assert resp.status_code == 403 + def test_create_memory_identity_less_caller_returns_400(self): + """ + A non-admin caller with neither user_id nor team_id would produce an + orphan row unreachable by the visibility filter. Reject up front. + """ + client = _make_client(UserAPIKeyAuth(api_key="sk-anon")) + with _patch_prisma(self.prisma): + resp = client.post("/v1/memory", json={"key": "notes", "value": "x"}) + assert resp.status_code == 400 + + def test_put_memory_admin_can_bootstrap_foreign_scope(self): + """PUT-create should mirror POST's admin scope override.""" + client = _make_client(_admin_auth()) + with _patch_prisma(self.prisma): + resp = client.put( + "/v1/memory/notes", + json={"value": "x", "user_id": "some-user", "team_id": "some-team"}, + ) + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["user_id"] == "some-user" + assert body["team_id"] == "some-team" + + def test_put_memory_race_returns_update_on_unique_violation(self): + """ + Simulate the check-then-create race: _find_memory_for_caller says the + row doesn't exist, then the create call gets a unique-constraint + violation (a concurrent writer beat us). The handler should re-read + and fall through to an update instead of surfacing a 500. + """ + table = self.prisma.db.litellm_memorytable + + original_create = table.create + original_find_many = table.find_many + pre_create_calls = {"n": 0} + + async def racing_create(data): + # On the very first create call we issue during the upsert, pretend + # another writer inserted the row just before us. + pre_create_calls["n"] += 1 + if pre_create_calls["n"] == 1: + table.rows.append( + _make_row( + memory_id="m-race", + key=data["key"], + value="from-other-writer", + user_id="user-a", + team_id="team-a", + ) + ) + raise Exception("UniqueViolation: duplicate key (raced)") + return await original_create(data) + + table.create = racing_create # type: ignore[assignment] + + client = _make_client(_user_auth("user-a", "team-a")) + with _patch_prisma(self.prisma): + resp = client.put("/v1/memory/notes", json={"value": "mine"}) + + assert resp.status_code == 200, resp.text + body = resp.json() + # Update path kicked in: our value replaced the racer's. + assert body["value"] == "mine" + # Only one row exists (the one the racer inserted, now updated). + assert len(table.rows) == 1 + + # Restore the fake's methods. + table.create = original_create # type: ignore[assignment] + table.find_many = original_find_many # type: ignore[assignment] + def test_admin_can_set_any_scope(self): client = _make_client(_admin_auth()) with _patch_prisma(self.prisma): diff --git a/ui/litellm-dashboard/src/components/MemoryView/MemoryView.tsx b/ui/litellm-dashboard/src/components/MemoryView/MemoryView.tsx index b5b51884e52..b6adcd178be 100644 --- a/ui/litellm-dashboard/src/components/MemoryView/MemoryView.tsx +++ b/ui/litellm-dashboard/src/components/MemoryView/MemoryView.tsx @@ -58,12 +58,20 @@ function formatTimestamp(ts?: string): string { } } +const PAGE_SIZE = 50; + export const MemoryView: React.FC = ({ accessToken }) => { const [searchInput, setSearchInput] = useState(""); const [appliedSearch, setAppliedSearch] = useState(""); const [detailRow, setDetailRow] = useState(null); const [editRow, setEditRow] = useState(null); const [isCreateOpen, setIsCreateOpen] = useState(false); + const [currentPage, setCurrentPage] = useState(1); + + // Reset to page 1 whenever the filter changes. + React.useEffect(() => { + setCurrentPage(1); + }, [appliedSearch]); const { data, @@ -71,20 +79,22 @@ export const MemoryView: React.FC = ({ accessToken }) => { isFetching, refetch, } = useQuery({ - queryKey: ["memoryList", appliedSearch], + queryKey: ["memoryList", appliedSearch, currentPage], queryFn: () => { if (!accessToken) throw new Error("Access token required"); // Prefix search matches the Redis-style mental model (namespace scan): // typing "user:" finds "user:profile", "user:prefs", etc. return fetchMemoryList(accessToken, { keyPrefix: appliedSearch || undefined, - pageSize: 200, + page: currentPage, + pageSize: PAGE_SIZE, }); }, enabled: !!accessToken, }); const rows = useMemo(() => data?.memories ?? [], [data]); + const total = data?.total ?? 0; const handleDelete = async (row: MemoryRow) => { Modal.confirm({ @@ -325,7 +335,18 @@ export const MemoryView: React.FC = ({ accessToken }) => { loading={isLoading} dataSource={rows} columns={columns} - pagination={{ pageSize: 20, showSizeChanger: true }} + // Server-side pagination: we fetch one page at a time so we never + // silently truncate large stores. `total` drives the page count; + // changing page/pageSize retriggers the query via `currentPage`. + pagination={{ + current: currentPage, + pageSize: PAGE_SIZE, + total, + showSizeChanger: false, + showTotal: (n, range) => + `${range[0]}–${range[1]} of ${n}`, + onChange: (page) => setCurrentPage(page), + }} locale={{ emptyText: ( = ({ accessToken }) => { ), }} /> - {data?.total !== undefined && ( - - Showing {rows.length} of {data.total} - - )}