Cache memory enrollment and coordinate continuation cleanup

litellm/proxy/memory/continuation.py

@@ -1,11 +1,13 @@
 from collections.abc import Mapping
 from datetime import datetime, timedelta, timezone
+from time import monotonic
 from types import MappingProxyType, SimpleNamespace
 from typing import Final
 
 from fastapi import HTTPException
 from pydantic import BaseModel, ConfigDict
 
+from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
 from litellm.proxy.memory.policy import memory_digest, memory_primary_client
 from litellm.proxy.memory.store import MemoryStore
 from litellm.repositories.table_repositories import MemoryContinuationRepository
@@ -14,16 +16,31 @@ from litellm.repositories.unit_of_work import prisma_transaction
 _MAX_PATCH_BYTES: Final = 1024 * 1024
 _MAX_PATCHES: Final = 256
 _MAX_NAMESPACE_BYTES: Final = 32 * 1024 * 1024
+MEMORY_CLEANUP_INTERVAL_SECONDS: Final = 3600
 
 
-async def cleanup_memory_continuations(prisma_client: object) -> None:
-    async with prisma_transaction(memory_primary_client(prisma_client)) as transaction:
-        await transaction.execute_raw(
-            'DELETE FROM "LiteLLM_MemoryContinuation" WHERE id IN '
-            '(SELECT id FROM "LiteLLM_MemoryContinuation" WHERE expires_at <= $1::timestamp '
-            "ORDER BY expires_at LIMIT 1000 FOR UPDATE SKIP LOCKED)",
-            datetime.now(timezone.utc),
+async def cleanup_memory_continuations(prisma_client: object, pod_lock_manager: PodLockManager) -> None:
+    if (
+        await pod_lock_manager.acquire_lock(
+            "memory_continuation_cleanup", ttl=MEMORY_CLEANUP_INTERVAL_SECONDS - 60, allow_reentrant=False
         )
+        is False
+    ):
+        return
+    deadline: Final = monotonic() + 10
+    for _ in range(100):
+        if monotonic() >= deadline:
+            return
+        async with prisma_transaction(memory_primary_client(prisma_client)) as transaction:
+            await transaction.execute_raw("SET LOCAL statement_timeout = '5s'")
+            deleted = await transaction.execute_raw(
+                'DELETE FROM "LiteLLM_MemoryContinuation" WHERE id IN '
+                '(SELECT id FROM "LiteLLM_MemoryContinuation" WHERE expires_at <= $1::timestamp '
+                "ORDER BY expires_at LIMIT 1000 FOR UPDATE SKIP LOCKED)",
+                datetime.now(timezone.utc),
+            )
+        if deleted < 1000:
+            return
 
 
 class MemoryContinuation(BaseModel):

litellm/proxy/memory/gateway.py

@@ -42,7 +42,7 @@ from litellm.proxy.memory.knowledge import (
 )
 from litellm.proxy.memory.policy import (
     MemoryIdentity,
-    gateway_memory_is_configured,
+    gateway_memory_is_enabled,
     resolve_memory_access,
 )
 from litellm.proxy.memory.store import MemoryStore
@@ -484,7 +484,7 @@ async def gateway_memory_store(auth: UserAPIKeyAuth) -> MemoryStore | None:
     if not identity.user_id and not identity.key_id:
         return None
     try:
-        if not await gateway_memory_is_configured(prisma_client, user_api_key_cache):
+        if not await gateway_memory_is_enabled(prisma_client, user_api_key_cache, identity):
             return None
         access: Final = await resolve_memory_access(prisma_client, identity)
         required_tools: Final = frozenset(str(function["name"]) for function in memory_functions(access))

litellm/proxy/memory/management.py

@@ -138,7 +138,7 @@ async def named_entries(store: MemoryStore, entries: tuple[MemoryEntry, ...]) ->
         entry.model_copy(
             update=MappingProxyType(
                 {
-                    "actor_name": names.get(entry.actor or ""),
+                    "actor_name": names.get(entry.actor or "", entry.actor_name),
                     "team_name": team_names.get(entry.team_id or ""),
                 }
             )

litellm/proxy/memory/policy.py

@@ -16,7 +16,7 @@ from litellm.repositories.team_repository import TeamRepository
 from litellm.repositories.user_repository import UserRepository
 from litellm.types.memory_v2 import MemorySettings, MemoryStatus
 
-_CONFIGURED_CACHE_KEY: Final = "litellm:memory_v2:configured"
+_SETTINGS_CACHE_KEY: Final = "litellm:memory_v2:settings"
 MEMORY_CONFIG_PARAM: Final = "memory_v2"
 
 
@@ -25,24 +25,23 @@ async def memory_settings(prisma_client: object) -> MemorySettings:
     return MemorySettings.model_validate(row.param_value) if row is not None else MemorySettings()
 
 
-async def gateway_memory_is_configured(prisma_client: object, cache: DualCache) -> bool:
-    cached: Final = await cache.async_get_cache(key=_CONFIGURED_CACHE_KEY)
-    if cached is True:
+async def gateway_memory_is_enabled(prisma_client: object, cache: DualCache, identity: "MemoryIdentity") -> bool:
+    cached: Final = await cache.async_get_cache(key=_SETTINGS_CACHE_KEY, ttl=30)
+    if cached is not None and MemoryAccess(identity, MemorySettings.model_validate(cached)).active:
         return True
     redis_cache: Final = cache.redis_cache or coordination_redis_cache()
     shared_cache: Final = DualCache(redis_cache=redis_cache) if redis_cache is not None else None
-    if cached is False:
+    if cached is not None:
         if shared_cache is None:
             return False
-        shared: Final = await shared_cache.async_get_cache(key=_CONFIGURED_CACHE_KEY)
-        if shared is False:
-            return False
+        shared: Final = await shared_cache.async_get_cache(key=_SETTINGS_CACHE_KEY, ttl=30)
+        if shared is not None:
+            return MemoryAccess(identity, MemorySettings.model_validate(shared)).active
     settings: Final = await memory_settings(prisma_client)
-    configured: Final = settings.enabled or settings.read.enabled
-    await cache.async_set_cache(key=_CONFIGURED_CACHE_KEY, value=configured, ttl=30)
+    await cache.async_set_cache(key=_SETTINGS_CACHE_KEY, value=settings.model_dump(mode="json"), ttl=30)
     if shared_cache is not None and cache.redis_cache is None:
-        await shared_cache.async_set_cache(key=_CONFIGURED_CACHE_KEY, value=configured, ttl=30)
-    return configured
+        await shared_cache.async_set_cache(key=_SETTINGS_CACHE_KEY, value=settings.model_dump(mode="json"), ttl=30)
+    return MemoryAccess(identity, settings).active
 
 
 async def invalidate_memory_configuration() -> None:
@@ -52,7 +51,7 @@ async def invalidate_memory_configuration() -> None:
         in_memory_cache=user_api_key_cache.in_memory_cache,
         redis_cache=user_api_key_cache.redis_cache or coordination_redis_cache(),
     )
-    await evict_and_broadcast(cache_keys=(_CONFIGURED_CACHE_KEY,), user_api_key_cache=cache)
+    await evict_and_broadcast(cache_keys=(_SETTINGS_CACHE_KEY,), user_api_key_cache=cache)
 
 
 def memory_primary_client(prisma_client: object) -> WriterPinnedClient:

litellm/proxy/memory/store.py

@@ -38,7 +38,8 @@ def memory_entry(row: "LiteLLM_MemoryTable") -> MemoryEntry:
             "evidence": evidence if isinstance(evidence, str) else "",
             "updated_at": row.updated_at,
             "created_at": row.created_at,
-            "actor": row.created_by,
+            "actor": row.created_by if row.user_id is not None else None,
+            "actor_name": "Service key" if row.user_id is None else None,
             "user_id": row.user_id,
             "team_id": row.team_id,
             **{  # mutable-ok: Prisma requires native JSON.

litellm/proxy/proxy_server.py

@@ -10251,13 +10251,13 @@ class ProxyStartupEvent:
         await cls._initialize_expired_ui_session_key_cleanup_background_job(scheduler=scheduler)
 
         if prisma_client is not None:
-            from litellm.proxy.memory.continuation import cleanup_memory_continuations
+            from litellm.proxy.memory.continuation import MEMORY_CLEANUP_INTERVAL_SECONDS, cleanup_memory_continuations
 
             scheduler.add_job(
                 cleanup_memory_continuations,
                 "interval",
-                seconds=60,
-                args=(prisma_client,),
+                seconds=MEMORY_CLEANUP_INTERVAL_SECONDS,
+                args=(prisma_client, proxy_logging_obj.db_spend_update_writer.pod_lock_manager),
                 id="memory_continuation_cleanup",
                 max_instances=1,
                 coalesce=True,

tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py

@@ -517,32 +517,33 @@ async def test_replica_lag_cannot_authorize_memory_after_primary_revocation(
 @pytest.mark.asyncio
 async def test_unconfigured_gate_caches_presence_without_caching_authorization(prisma_edge: MagicMock) -> None:
     from litellm.caching.caching import DualCache
-    from litellm.proxy.memory.policy import gateway_memory_is_configured
+    from litellm.proxy.memory.policy import gateway_memory_is_enabled
 
     cache = DualCache()
     config = prisma_edge.db.litellm_config.find_unique
     config.return_value = None
-    assert not await gateway_memory_is_configured(prisma_edge, cache)
-    assert not await gateway_memory_is_configured(prisma_edge, cache)
+    assert not await gateway_memory_is_enabled(prisma_edge, cache, _IDENTITY)
+    assert not await gateway_memory_is_enabled(prisma_edge, cache, _IDENTITY)
     config.assert_awaited_once()
     assert config.call_args.kwargs == {"where": {"param_name": "memory_v2"}}
     enabled_cache = DualCache()
     config.return_value = SimpleNamespace(param_value=_SETTINGS.model_dump())
-    assert await gateway_memory_is_configured(prisma_edge, enabled_cache)
+    assert await gateway_memory_is_enabled(prisma_edge, enabled_cache, _IDENTITY)
     config.return_value = None
-    assert await gateway_memory_is_configured(prisma_edge, enabled_cache)
+    assert await gateway_memory_is_enabled(prisma_edge, enabled_cache, _IDENTITY)
     assert not (await resolve_memory_access(prisma_edge, _IDENTITY)).active
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("share_auth_cache", [False, True])
+@pytest.mark.parametrize("other_user_enrolled", [False, True])
 async def test_backend_activation_invalidates_a_gateway_negative_hint_without_pubsub(
-    prisma_edge: MagicMock, share_auth_cache: bool
+    prisma_edge: MagicMock, share_auth_cache: bool, other_user_enrolled: bool
 ) -> None:
     from unittest.mock import patch
 
     from litellm.caching.caching import DualCache
-    from litellm.proxy.memory.policy import gateway_memory_is_configured, invalidate_memory_configuration
+    from litellm.proxy.memory.policy import gateway_memory_is_enabled, invalidate_memory_configuration
 
     shared = {}
 
@@ -566,13 +567,15 @@ async def test_backend_activation_invalidates_a_gateway_negative_hint_without_pu
     with patch.multiple(  # test-quality-ok: Inject external worker caches and Redis; run real invalidation.
         "litellm.proxy.proxy_server", user_api_key_cache=backend_cache, redis_usage_cache=redis
     ):
-        config.return_value = None
-        assert not await gateway_memory_is_configured(prisma_edge, gateway_cache)
-        assert not await gateway_memory_is_configured(prisma_edge, gateway_cache)
+        config.return_value = SimpleNamespace(
+            param_value=MemorySettings(enabled=other_user_enrolled, everyone=False, user_ids=("other",)).model_dump()
+        )
+        assert not await gateway_memory_is_enabled(prisma_edge, gateway_cache, _IDENTITY)
+        assert not await gateway_memory_is_enabled(prisma_edge, gateway_cache, _IDENTITY)
         config.assert_awaited_once()
         config.return_value = SimpleNamespace(param_value=_SETTINGS.model_dump())
         await invalidate_memory_configuration()
-        assert await gateway_memory_is_configured(prisma_edge, gateway_cache)
+        assert await gateway_memory_is_enabled(prisma_edge, gateway_cache, _IDENTITY)
         assert config.await_count == 2
 
 
@@ -582,7 +585,7 @@ async def test_redis_circuit_breaker_falls_back_to_primary_configuration(prisma_
 
     from litellm.caching.caching import DualCache
     from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
-    from litellm.proxy.memory.policy import gateway_memory_is_configured, invalidate_memory_configuration
+    from litellm.proxy.memory.policy import gateway_memory_is_enabled, invalidate_memory_configuration
 
     redis = MagicMock(
         async_get_cache=AsyncMock(side_effect=RedisCircuitBreakerOpenError("open")),
@@ -594,14 +597,68 @@ async def test_redis_circuit_breaker_falls_back_to_primary_configuration(prisma_
     with patch.multiple(  # test-quality-ok: Inject external Redis failure and local worker cache; exercise real fallback.
         "litellm.proxy.proxy_server", user_api_key_cache=cache, redis_usage_cache=redis
     ):
-        assert not await gateway_memory_is_configured(prisma_edge, cache)
+        assert not await gateway_memory_is_enabled(prisma_edge, cache, _IDENTITY)
         prisma_edge.db.litellm_config.find_unique.return_value = SimpleNamespace(param_value=_SETTINGS.model_dump())
-        assert await gateway_memory_is_configured(prisma_edge, cache)
+        assert await gateway_memory_is_enabled(prisma_edge, cache, _IDENTITY)
         await invalidate_memory_configuration()
     assert prisma_edge.db.litellm_config.find_unique.await_count == 2
     redis.async_get_cache.assert_awaited_once()
 
 
+@pytest.mark.asyncio
+async def test_cleanup_coordinates_workers_and_keeps_completed_window_locked(prisma_edge: MagicMock) -> None:
+    from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
+    from litellm.proxy.memory.continuation import cleanup_memory_continuations
+
+    shared = {}
+    now = 0
+
+    async def set_value(key, value, *, nx, ttl):
+        assert nx
+        if key in shared and shared[key][1] > now:
+            return False
+        shared[key] = (value, now + ttl)
+        return True
+
+    async def get(key):
+        return shared[key][0] if key in shared and shared[key][1] > now else None
+
+    redis = MagicMock(async_set_cache=AsyncMock(side_effect=set_value), async_get_cache=AsyncMock(side_effect=get))
+    first = PodLockManager(redis)
+    second = PodLockManager(redis)
+    prisma_edge.db.execute_raw.return_value = 0
+    await cleanup_memory_continuations(prisma_edge, first)
+    transactions = prisma_edge.db.tx.call_count
+    assert transactions == 1
+    await cleanup_memory_continuations(prisma_edge, second)
+    await cleanup_memory_continuations(prisma_edge, first)
+    assert prisma_edge.db.tx.call_count == transactions
+    now = 3599
+    await cleanup_memory_continuations(prisma_edge, second)
+    assert prisma_edge.db.tx.call_count == transactions + 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("full_batches", [0, 2, 100])
+async def test_cleanup_without_redis_drains_bounded_batches(prisma_edge: MagicMock, full_batches: int) -> None:
+    from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
+    from litellm.proxy.memory.continuation import cleanup_memory_continuations
+
+    sizes = iter([1000] * full_batches + [17])
+
+    async def execute(query, *params):
+        if query.startswith("DELETE"):
+            assert "expires_at <= $1" in query and "LIMIT 1000 FOR UPDATE SKIP LOCKED" in query
+            assert params[0] <= datetime.now(timezone.utc)
+            return next(sizes)
+        assert query == "SET LOCAL statement_timeout = '5s'"
+        return 0
+
+    prisma_edge.db.execute_raw.side_effect = execute
+    await cleanup_memory_continuations(prisma_edge, PodLockManager())
+    assert prisma_edge.db.tx.call_count == min(full_batches + 1, 100)
+
+
 @pytest.mark.asyncio
 async def test_full_scope_blocks_creation_but_permits_correction_and_reclaimed_capacity(prisma_edge: MagicMock) -> None:
     table = prisma_edge.db.litellm_memorytable

tests/test_litellm/proxy/memory/test_memory_v2_management.py

@@ -298,6 +298,51 @@ async def test_service_key_write_ownership_never_matches_all_unowned_rows(databa
     assert store.access.visible_rows(write=True)["OR"] == [{"owner_key_id": "a" * 64, "user_id": None}]
 
 
+@pytest.mark.asyncio
+async def test_service_key_hash_is_private_in_entries_and_agent_results(database: MagicMock) -> None:
+    configure(database, read={"enabled": True})
+    database.db.litellm_teamtable.find_many.return_value = [team(permissions=("/memory/v2/entries",))]
+    service_row = row(user_id=None, created_by="b" * 64, owner_key_id="b" * 64)
+    table = database.db.litellm_memorytable
+    table.find_first.return_value = service_row
+    table.find_many.return_value = [service_row]
+    table.create.return_value = service_row
+    reader = await management.memory_store(auth())
+    entries = await management.list_entries(query="", limit=20, offset=0, auth=auth())
+    named = await management.read_entry("entry", auth=auth())
+    recalled, _ = await reader.recall(MemoryRecallRequest(query=""))
+    captured = await management.capture_entry(_CAPTURE, auth(None).model_copy(update={"token": "b" * 64}))
+    for entry in (*entries, named, recalled[0][0], captured):
+        assert entry.actor is None and entry.actor_name == "Service key"
+        assert "b" * 64 not in entry.model_dump_json()
+    assert not named.can_edit and captured.can_edit
+    assert table.create.call_args.kwargs["data"]["created_by"] == "b" * 64
+    database.db.litellm_usertable.find_many.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("read_only_enrollment", [False, True])
+async def test_unenrolled_requests_skip_user_and_team_reads_and_settings_changes_invalidate(
+    database: MagicMock, read_only_enrollment: bool
+) -> None:
+    enrollment = {"enabled": True, "everyone": False, "user_ids": ["other"]}
+    configure(database, **({"enabled": False, "read": enrollment} if read_only_enrollment else enrollment))
+    for _ in range(10):
+        assert await gateway_memory_store(auth()) is None
+    database.db.litellm_config.find_unique.assert_awaited_once()
+    database.db.litellm_usertable.find_unique.assert_not_awaited()
+    database.db.litellm_teamtable.find_many.assert_not_awaited()
+    enrolled = await gateway_memory_store(auth("other"))
+    assert enrolled is not None and enrolled.access.active
+    settings = MemorySettings(enabled=True, read=MemoryEnrollment(enabled=True))
+    await management.set_settings(settings, auth(role=LitellmUserRoles.PROXY_ADMIN))
+    configure(database, **settings.model_dump())
+    newly_enrolled = await gateway_memory_store(auth())
+    assert newly_enrolled is not None and newly_enrolled.access.active
+    configure(database, enabled=False)
+    assert await gateway_memory_store(auth()) is None
+
+
 @pytest.mark.asyncio
 async def test_revocation_blocks_existing_store_and_private_continuation(database: MagicMock) -> None:
     configure(database)