From 396b18b74b97ee1135d1f4480549c625c1531df0 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 15 Sep 2026 16:34:20 -0700 Subject: [PATCH] Cache memory enrollment and coordinate continuation cleanup --- litellm/proxy/memory/continuation.py | 31 +++++-- litellm/proxy/memory/gateway.py | 4 +- litellm/proxy/memory/management.py | 2 +- litellm/proxy/memory/policy.py | 25 +++--- litellm/proxy/memory/store.py | 3 +- litellm/proxy/proxy_server.py | 6 +- .../proxy/memory/test_memory_v2_boundaries.py | 85 ++++++++++++++++--- .../proxy/memory/test_memory_v2_management.py | 45 ++++++++++ 8 files changed, 160 insertions(+), 41 deletions(-) diff --git a/litellm/proxy/memory/continuation.py b/litellm/proxy/memory/continuation.py index e405b8df059..be0efbfb04a 100644 --- a/litellm/proxy/memory/continuation.py +++ b/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): diff --git a/litellm/proxy/memory/gateway.py b/litellm/proxy/memory/gateway.py index ae9c4332646..77a089d0330 100644 --- a/litellm/proxy/memory/gateway.py +++ b/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)) diff --git a/litellm/proxy/memory/management.py b/litellm/proxy/memory/management.py index cc503be64e5..ab840993438 100644 --- a/litellm/proxy/memory/management.py +++ b/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 ""), } ) diff --git a/litellm/proxy/memory/policy.py b/litellm/proxy/memory/policy.py index 1dab9109c86..287a73c839b 100644 --- a/litellm/proxy/memory/policy.py +++ b/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: diff --git a/litellm/proxy/memory/store.py b/litellm/proxy/memory/store.py index 0fd2b62a601..715a28fe397 100644 --- a/litellm/proxy/memory/store.py +++ b/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. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e0bdba86fd6..01df79d1257 100644 --- a/litellm/proxy/proxy_server.py +++ b/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, diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py index 74423796e0e..7847b647187 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py +++ b/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 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 9525333600d..55840bef358 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_management.py +++ b/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)