From 3b9445eeb2bda64a51090128e23614c66ed546bd Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 23 Apr 2026 11:15:19 -0700 Subject: [PATCH] fix(memory): make key globally unique, 409 on any duplicate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Switches from the compound unique `(key, user_id, team_id)` to a simple `key @unique`. The compound form silently allowed duplicates when user_id or team_id was NULL (Postgres treats each NULL as distinct), so callers could POST the same key repeatedly. Globally-unique key means one row per key, period — any duplicate create → 409. - schema.prisma (×3): `key String @unique`, drop `@@unique(...)`. - initial add_memory_table migration: unique index on (key) only. - Remove the now-unused follow-up NULLS NOT DISTINCT migration. - Endpoint error message simplified ("already exists" — no "for this scope"). - Test fake's create() now enforces global key uniqueness. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../migration.sql | 5 +-- .../migration.sql | 34 ------------------- .../litellm_proxy_extras/schema.prisma | 3 +- litellm/proxy/memory/memory_endpoints.py | 4 +-- litellm/proxy/schema.prisma | 3 +- schema.prisma | 3 +- .../proxy/memory/test_memory_endpoints.py | 7 ++-- 7 files changed, 10 insertions(+), 49 deletions(-) delete mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260423090000_memory_unique_nulls_not_distinct/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260421120000_add_memory_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260421120000_add_memory_table/migration.sql index 76db89203c3..8bba29d53af 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260421120000_add_memory_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260421120000_add_memory_table/migration.sql @@ -14,8 +14,9 @@ CREATE TABLE IF NOT EXISTS "LiteLLM_MemoryTable" ( CONSTRAINT "LiteLLM_MemoryTable_pkey" PRIMARY KEY ("memory_id") ); --- CreateIndex -CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_key_user_id_team_id_key" ON "LiteLLM_MemoryTable"("key", "user_id", "team_id"); +-- CreateIndex (key is globally unique — one row per key, period) +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_key_key" + ON "LiteLLM_MemoryTable"("key"); -- CreateIndex CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_user_id_idx" ON "LiteLLM_MemoryTable"("user_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260423090000_memory_unique_nulls_not_distinct/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260423090000_memory_unique_nulls_not_distinct/migration.sql deleted file mode 100644 index 3846d3961b6..00000000000 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260423090000_memory_unique_nulls_not_distinct/migration.sql +++ /dev/null @@ -1,34 +0,0 @@ --- Fix duplicate-key issue on LiteLLM_MemoryTable. --- --- Problem: the unique constraint `(key, user_id, team_id)` does not block --- duplicates when `user_id` or `team_id` is NULL, because by default --- Postgres treats each NULL as distinct (ANSI SQL semantics). Callers with --- no team_id could POST the same key repeatedly and get multiple rows. --- --- Fix (Postgres 15+): recreate the unique index with NULLS NOT DISTINCT so --- NULL values are treated as equal for uniqueness checks. --- --- Note: Prisma has no schema syntax for NULLS NOT DISTINCT (as of writing), --- so this lives as a raw SQL migration. Prisma introspection will still see --- the index as uniquely covering `(key, user_id, team_id)`, matching the --- `@@unique` in schema.prisma. - --- 1. Deduplicate existing rows: keep only the most-recently-updated row per --- (key, user_id, team_id), using IS NOT DISTINCT FROM so NULL == NULL. -DELETE FROM "LiteLLM_MemoryTable" a -USING "LiteLLM_MemoryTable" b -WHERE a.memory_id <> b.memory_id - AND a.key = b.key - AND a.user_id IS NOT DISTINCT FROM b.user_id - AND a.team_id IS NOT DISTINCT FROM b.team_id - AND ( - a.updated_at < b.updated_at - OR (a.updated_at = b.updated_at AND a.memory_id < b.memory_id) - ); - --- 2. Drop the old (NULL-distinct) unique index. -DROP INDEX IF EXISTS "LiteLLM_MemoryTable_key_user_id_team_id_key"; - --- 3. Recreate with NULLS NOT DISTINCT. -CREATE UNIQUE INDEX "LiteLLM_MemoryTable_key_user_id_team_id_key" - ON "LiteLLM_MemoryTable"("key", "user_id", "team_id") NULLS NOT DISTINCT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 344cc16d76e..f23956d5c5b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1229,7 +1229,7 @@ model LiteLLM_ClaudeCodePluginTable { // `metadata` is an optional JSON envelope for structured tags without schema changes. model LiteLLM_MemoryTable { memory_id String @id @default(uuid()) - key String + key String @unique value String metadata Json? user_id String? @@ -1239,7 +1239,6 @@ model LiteLLM_MemoryTable { updated_at DateTime @default(now()) @updatedAt updated_by String? - @@unique([key, user_id, team_id]) @@index([user_id]) @@index([team_id]) } diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index 1bb8b808977..bdc73046c8c 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -138,12 +138,12 @@ async def create_memory( try: row = await prisma_client.db.litellm_memorytable.create(data=create_data) except Exception as e: - # Unique constraint (key, user_id, team_id) → 409. + # Key is globally unique. Any duplicate → 409. msg = str(e) if "Unique" in msg or "unique" in msg or "UniqueViolation" in msg: raise HTTPException( status_code=409, - detail=f"Memory with key '{body.key}' already exists for this scope.", + detail=f"Memory with key '{body.key}' already exists.", ) verbose_proxy_logger.exception("Error creating memory: %s", e) raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 344cc16d76e..f23956d5c5b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1229,7 +1229,7 @@ model LiteLLM_ClaudeCodePluginTable { // `metadata` is an optional JSON envelope for structured tags without schema changes. model LiteLLM_MemoryTable { memory_id String @id @default(uuid()) - key String + key String @unique value String metadata Json? user_id String? @@ -1239,7 +1239,6 @@ model LiteLLM_MemoryTable { updated_at DateTime @default(now()) @updatedAt updated_by String? - @@unique([key, user_id, team_id]) @@index([user_id]) @@index([team_id]) } diff --git a/schema.prisma b/schema.prisma index 344cc16d76e..f23956d5c5b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1229,7 +1229,7 @@ model LiteLLM_ClaudeCodePluginTable { // `metadata` is an optional JSON envelope for structured tags without schema changes. model LiteLLM_MemoryTable { memory_id String @id @default(uuid()) - key String + key String @unique value String metadata Json? user_id String? @@ -1239,7 +1239,6 @@ model LiteLLM_MemoryTable { updated_at DateTime @default(now()) @updatedAt updated_by String? - @@unique([key, user_id, team_id]) @@index([user_id]) @@index([team_id]) } diff --git a/tests/test_litellm/proxy/memory/test_memory_endpoints.py b/tests/test_litellm/proxy/memory/test_memory_endpoints.py index 2a9de6eac56..9733b5d9ec8 100644 --- a/tests/test_litellm/proxy/memory/test_memory_endpoints.py +++ b/tests/test_litellm/proxy/memory/test_memory_endpoints.py @@ -84,12 +84,9 @@ class _InMemoryMemoryTable: return [r for r in self.rows if self._matches(r, where)] async def create(self, data: Dict[str, Any]) -> MagicMock: + # Key is globally unique. for r in self.rows: - if ( - r.key == data["key"] - and r.user_id == data.get("user_id") - and r.team_id == data.get("team_id") - ): + if r.key == data["key"]: raise Exception("UniqueViolation: duplicate key") self._counter += 1 row = _make_row(