From cc1f0d4cf30963eb1984cd1e538f1c79e18f877e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 23 Apr 2026 08:24:07 -0700 Subject: [PATCH] feat(memory): add key_prefix filter + promote Memory to AI GATEWAY nav MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Backend: - GET /v1/memory now accepts `key_prefix` for Redis-style namespace scans (e.g. `?key_prefix=user:`). When both `key` and `key_prefix` are passed, `key_prefix` wins. - Prefix filter sits under the visibility filter in the Prisma where clause, so it can never leak rows across user/team scopes. - New tests: prefix match, and cross-scope isolation (another user's `user:*` rows must not appear in the caller's results). UI: - Memory moved from a Tools submenu to a top-level AI GATEWAY item (alongside Agents, MCP Servers, Skills) — it's an API primitive, not a tool-management surface. - Search box now drives prefix search, matching the Redis mental model ("type the namespace, see everything under it"). Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/proxy/memory/memory_endpoints.py | 14 +++- .../proxy/memory/test_memory_endpoints.py | 73 +++++++++++++++++++ .../src/components/MemoryView/MemoryView.tsx | 7 +- .../src/components/leftnav.tsx | 12 +-- .../src/components/networking.tsx | 15 +++- 5 files changed, 109 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index ff105c739cd..1bb8b808977 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -159,6 +159,13 @@ async def create_memory( ) async def list_memory( key: Optional[str] = Query(None, description="Filter by exact key match."), + key_prefix: Optional[str] = Query( + None, + description=( + "Filter by key prefix (Redis-style namespace scan). " + "Mutually exclusive with `key`; if both are provided, `key_prefix` wins." + ), + ), page: int = Query(1, ge=1), page_size: int = Query(50, ge=1, le=500), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -170,7 +177,12 @@ async def list_memory( vis = _visibility_filter(user_api_key_dict) if vis is not None: where.update(vis) - if key is not None: + # key_prefix takes precedence over exact `key` if both are passed. + # Both sit under the visibility filter (ANDed), so prefix scans can + # never leak across user/team scopes. + if key_prefix is not None: + where["key"] = {"startsWith": key_prefix} + elif key is not None: where["key"] = key try: diff --git a/tests/test_litellm/proxy/memory/test_memory_endpoints.py b/tests/test_litellm/proxy/memory/test_memory_endpoints.py index 4bbea02e32b..2a9de6eac56 100644 --- a/tests/test_litellm/proxy/memory/test_memory_endpoints.py +++ b/tests/test_litellm/proxy/memory/test_memory_endpoints.py @@ -59,6 +59,21 @@ class _InMemoryMemoryTable: if not any(self._matches(row, clause) for clause in v): return False continue + # Support Prisma-style filter dicts: {"startsWith": "..."}, etc. + if isinstance(v, dict): + actual = getattr(row, k, None) + if "startsWith" in v: + if not isinstance(actual, str) or not actual.startswith( + v["startsWith"] + ): + return False + continue + if "equals" in v: + if actual != v["equals"]: + return False + continue + # Unknown filter — fall back to inequality (treat as no match). + return False if getattr(row, k, None) != v: return False return True @@ -226,6 +241,64 @@ class TestMemoryEndpoints: assert keys == {"a", "c"} assert body["total"] == 2 + def test_list_memory_key_prefix_filter(self): + """key_prefix should do a prefix match (Redis-style namespace scan).""" + table = self.prisma.db.litellm_memorytable + table.rows.extend( + [ + _make_row( + memory_id="m1", key="user:profile", user_id="user-a", team_id=None + ), + _make_row( + memory_id="m2", key="user:prefs", user_id="user-a", team_id=None + ), + _make_row( + memory_id="m3", + key="project:context", + user_id="user-a", + team_id=None, + ), + ] + ) + client = _make_client(_user_auth("user-a", "team-a")) + with _patch_prisma(self.prisma): + resp = client.get("/v1/memory?key_prefix=user:") + assert resp.status_code == 200 + body = resp.json() + keys = {m["key"] for m in body["memories"]} + assert keys == {"user:profile", "user:prefs"} + assert body["total"] == 2 + + def test_list_memory_key_prefix_does_not_leak_across_scopes(self): + """ + Even if a prefix would match another user's keys, the visibility + filter must still scope results to the caller. + """ + table = self.prisma.db.litellm_memorytable + table.rows.extend( + [ + # Caller's own "user:*" rows — should be visible. + _make_row( + memory_id="m1", key="user:profile", user_id="user-a", team_id=None + ), + # Another user's "user:*" rows — must NOT leak. + _make_row( + memory_id="m2", key="user:secret", user_id="user-b", team_id=None + ), + _make_row( + memory_id="m3", key="user:token", user_id="user-b", team_id=None + ), + ] + ) + client = _make_client(_user_auth("user-a", "team-a")) + with _patch_prisma(self.prisma): + resp = client.get("/v1/memory?key_prefix=user:") + assert resp.status_code == 200 + body = resp.json() + keys = {m["key"] for m in body["memories"]} + assert keys == {"user:profile"} + assert body["total"] == 1 + def test_list_memory_admin_sees_all(self): table = self.prisma.db.litellm_memorytable table.rows.extend( diff --git a/ui/litellm-dashboard/src/components/MemoryView/MemoryView.tsx b/ui/litellm-dashboard/src/components/MemoryView/MemoryView.tsx index 21f09e1d4ef..b2f8c41436d 100644 --- a/ui/litellm-dashboard/src/components/MemoryView/MemoryView.tsx +++ b/ui/litellm-dashboard/src/components/MemoryView/MemoryView.tsx @@ -74,7 +74,8 @@ export const MemoryView: React.FC = ({ accessToken }) => { queryKey: ["memoryList", appliedSearch], queryFn: () => { if (!accessToken) throw new Error("Access token required"); - // Current API supports exact key match; treat empty search as list-all. + // 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, @@ -254,7 +255,7 @@ export const MemoryView: React.FC = ({ accessToken }) => { } value={searchInput} onChange={(e) => setSearchInput(e.target.value)} @@ -300,7 +301,7 @@ export const MemoryView: React.FC = ({ accessToken }) => { diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 272d348e3fd..c340b65496c 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -139,6 +139,12 @@ const menuGroups: MenuGroup[] = [ icon: , roles: all_admin_roles, }, + { + key: "memory", + page: "memory", + label: "Memory", + icon: , + }, { key: "guardrails", page: "guardrails", @@ -180,12 +186,6 @@ const menuGroups: MenuGroup[] = [ label: "Tool Policies", icon: , }, - { - key: "memory", - page: "memory", - label: "Memory", - icon: , - }, ], }, ], diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 1f08e0190f1..3355079f7bc 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -10011,11 +10011,22 @@ export interface MemoryListResponse { export const fetchMemoryList = async ( accessToken: string, - options: { keyPrefix?: string; page?: number; pageSize?: number } = {}, + options: { + key?: string; + keyPrefix?: string; + page?: number; + pageSize?: number; + } = {}, ): Promise => { const base = proxyBaseUrl ? `${proxyBaseUrl}/v1/memory` : `/v1/memory`; const params = new URLSearchParams(); - if (options.keyPrefix) params.append("key", options.keyPrefix); + // keyPrefix takes precedence — backend also does, but we omit `key` + // to keep the URL clean and intent obvious. + if (options.keyPrefix) { + params.append("key_prefix", options.keyPrefix); + } else if (options.key) { + params.append("key", options.key); + } if (options.page != null) params.append("page", String(options.page)); if (options.pageSize != null) params.append("page_size", String(options.pageSize));