From 7297247aab0094b688410fe1193dc7ddb180c92a Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 12 Sep 2026 02:51:40 -0700 Subject: [PATCH] fix(memory): invalidate worker hints and bound pilot validation --- deploy/memory-pilot/README.md | 3 + deploy/memory-pilot/pilot.py | 13 + litellm/proxy/memory/policy.py | 15 +- .../proxy/memory/test_memory_pilot.py | 54 ++++ .../proxy/memory/test_memory_v2_boundaries.py | 44 ++- .../proxy/memory/test_memory_v2_management.py | 254 ++++++++++++++++++ 6 files changed, 378 insertions(+), 5 deletions(-) create mode 100644 tests/test_litellm/proxy/memory/test_memory_pilot.py create mode 100644 tests/test_litellm/proxy/memory/test_memory_v2_management.py diff --git a/deploy/memory-pilot/README.md b/deploy/memory-pilot/README.md index 528d92a4ded..ab10d2f0e4a 100644 --- a/deploy/memory-pilot/README.md +++ b/deploy/memory-pilot/README.md @@ -83,6 +83,9 @@ correct, or delete entries in Memory; callers can use the self-service API. - Preparation stores durable facts supported by the conversation, then searches and reads relevant entries. Search is bounded keyword matching in Postgres. There is no vector database, extraction model, scheduler, or nightly process. +- On gateway/backend deployments without shared Redis, first-time activation + can take up to 30 seconds to reach another process. Policy revocation is + checked against the primary database before memory operations. - Stored references are untrusted data. They cannot grant API permissions or change the namespace derived from authentication. Current user corrections take precedence. Replacements require the current revision. diff --git a/deploy/memory-pilot/pilot.py b/deploy/memory-pilot/pilot.py index fc2083fc701..76cb00dff98 100644 --- a/deploy/memory-pilot/pilot.py +++ b/deploy/memory-pilot/pilot.py @@ -1,5 +1,6 @@ """An isolated office pilot that preserves upstream gateway credentials.""" +import asyncio import hashlib import os import secrets @@ -46,6 +47,7 @@ forward_credential: Final = ForwardCredential() class PilotGateway: def __init__(self, app: ASGIApp) -> None: self.app = app + self.validation_slots = asyncio.Semaphore(16) self.upstream = get_async_httpx_client( httpxSpecialProvider.PassThroughEndpoint, params={"timeout": 20, "client_alias": "memory-pilot-upstream"}, @@ -96,6 +98,15 @@ class PilotGateway: {"error": "Upstream keys can only use inference and their own memories"}, status_code=403 )(scope, receive, send) return + try: + await asyncio.wait_for(self.validation_slots.acquire(), timeout=0.05) + except TimeoutError: + await JSONResponse( + {"error": "Pilot credential validation is busy; retry shortly"}, + status_code=503, + headers={"Retry-After": "1"}, + )(scope, receive, send) + return try: models: Final = await self.upstream.get( _UPSTREAM + "/v1/models", headers={"Authorization": "Bearer " + credential} @@ -103,6 +114,8 @@ class PilotGateway: except httpx.HTTPError: await JSONResponse({"error": "Upstream gateway unavailable"}, status_code=503)(scope, receive, send) return + finally: + self.validation_slots.release() if models.is_error: await JSONResponse({"error": "Upstream gateway rejected this key"}, status_code=models.status_code)( scope, receive, send diff --git a/litellm/proxy/memory/policy.py b/litellm/proxy/memory/policy.py index 6ca18cc4909..666aa3ed948 100644 --- a/litellm/proxy/memory/policy.py +++ b/litellm/proxy/memory/policy.py @@ -7,6 +7,7 @@ from fastapi import HTTPException from litellm.caching.caching import DualCache from litellm.proxy._types import UI_TEAM_ID, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient from litellm.repositories.table_repositories import MemoryPolicyRepository, MemoryPreferenceRepository from litellm.types.memory_v2 import MemoryPolicy, MemoryScope, MemoryStatus @@ -22,8 +23,16 @@ async def gateway_memory_is_configured(prisma_client: object, cache: DualCache) without shared Redis when an administrator first enables memory. """ cached: Final = await cache.async_get_cache(key=_CONFIGURED_CACHE_KEY) - if isinstance(cached, bool): - return cached + if cached is True: + return True + if cached is False: + if cache.redis_cache is None: + return False + # A backend mutation evicts Redis, but another worker can still hold + # a negative local hint (Redis Cluster may not support pub/sub). + shared: Final = await cache.redis_cache.async_get_cache(key=_CONFIGURED_CACHE_KEY) + if shared is False: + return False rows: Final = await MemoryPolicyRepository(memory_primary_client(prisma_client)).table.find_many(take=1) configured: Final = bool(rows) await cache.async_set_cache(key=_CONFIGURED_CACHE_KEY, value=configured, ttl=30) @@ -33,7 +42,7 @@ async def gateway_memory_is_configured(prisma_client: object, cache: DualCache) async def invalidate_memory_configuration() -> None: from litellm.proxy.proxy_server import user_api_key_cache - await user_api_key_cache.async_delete_cache(key=_CONFIGURED_CACHE_KEY) + await evict_and_broadcast(cache_keys=(_CONFIGURED_CACHE_KEY,), user_api_key_cache=user_api_key_cache) def memory_primary_client(prisma_client: object) -> WriterPinnedClient: diff --git a/tests/test_litellm/proxy/memory/test_memory_pilot.py b/tests/test_litellm/proxy/memory/test_memory_pilot.py new file mode 100644 index 00000000000..bb2e8d29499 --- /dev/null +++ b/tests/test_litellm/proxy/memory/test_memory_pilot.py @@ -0,0 +1,54 @@ +"""Bound unauthenticated upstream validation without replacing gateway authentication.""" + +import asyncio +import importlib.util +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest +from starlette.applications import Starlette + + +@pytest.mark.asyncio +async def test_upstream_validation_is_bounded_and_slots_release_after_failure(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("UPSTREAM_LITELLM_BASE_URL", "https://upstream.example.invalid") + filename = Path(__file__).resolve().parents[4] / "deploy" / "memory-pilot" / "pilot.py" + spec = importlib.util.spec_from_file_location("memory_pilot_test", filename) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + gateway = module.PilotGateway(Starlette()) + entered = asyncio.Event() + release = asyncio.Event() + count = 0 + + async def upstream_get(*args: object, **kwargs: object) -> httpx.Response: + nonlocal count + count += 1 + if count == 16: + entered.set() + await release.wait() + raise httpx.ConnectError("unavailable") + + upstream = MagicMock(get=AsyncMock(side_effect=upstream_get)) + gateway.upstream = upstream + database = MagicMock() + database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + with patch.multiple( # test-quality-ok: Inject external database/config; exercise real ASGI admission. + "litellm.proxy.proxy_server", prisma_client=database, master_key="local-admin" + ): + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=gateway), base_url="http://pilot") as client: + pending = [ + asyncio.create_task(client.get("/v1/models", headers={"Authorization": f"Bearer sk-invalid-{i}"})) + for i in range(16) + ] + await asyncio.wait_for(entered.wait(), timeout=2) + refused = await client.get("/v1/models", headers={"Authorization": "Bearer sk-overload"}) + assert refused.status_code == 503 and refused.headers["retry-after"] == "1" + assert upstream.get.await_count == 16 + release.set() + assert all(response.status_code == 503 for response in await asyncio.gather(*pending)) + again = await client.get("/v1/models", headers={"Authorization": "Bearer sk-next"}) + assert again.status_code == 503 and "unavailable" in again.text + assert upstream.get.await_count == 17 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 4e2481e2c86..c1ffb02f82d 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py +++ b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py @@ -351,8 +351,9 @@ async def test_gateway_rounds_keep_separate_limiter_contexts_and_original_client ) auth = UserAPIKeyAuth(token="a" * 64, user_id="owner", team_id="team", project_id="project", org_id="org") before = get_request_stash() - # test-quality-ok: inject only database, cache, and ASGI provider edges; the memory and limiter code runs unchanged. - with patch.multiple(proxy_server, app=provider, prisma_client=prisma_edge, user_api_key_cache=DualCache()): + with patch.multiple( # test-quality-ok: Inject database/cache/ASGI provider edges; run real memory and limiter code. + proxy_server, app=provider, prisma_client=prisma_edge, user_api_key_cache=DualCache() + ): prepared = await prepare_gateway_memory(original, request, auth, "anthropic_messages") release.set() owners = await asyncio.gather(*deferred) @@ -364,3 +365,42 @@ async def test_gateway_rounds_keep_separate_limiter_contexts_and_original_client assert prepared["tools"] == original["tools"] and prepared["tool_choice"] == original["tool_choice"] assert prepared["stream"] is True and prepared["max_tokens"] == 1 assert "8347" in json.dumps(prepared) and "8347" not in json.dumps(original) + + +@pytest.mark.asyncio +async def test_backend_activation_invalidates_a_gateway_negative_hint_without_pubsub(prisma_edge: MagicMock) -> 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 + + shared = {} + + async def get(key, **kwargs): + return shared.get(key) + + async def set_value(key, value, **kwargs): + shared[key] = value + + async def delete(key, **kwargs): + shared.pop(key, None) + + redis = MagicMock( + async_get_cache=AsyncMock(side_effect=get), + async_set_cache=AsyncMock(side_effect=set_value), + async_delete_cache=AsyncMock(side_effect=delete), + ) + gateway_cache = DualCache(redis_cache=redis) + backend_cache = DualCache(redis_cache=redis) + policies = prisma_edge.db.litellm_memorypolicy.find_many + policies.return_value = [] + assert not await gateway_memory_is_configured(prisma_edge, gateway_cache) + assert not await gateway_memory_is_configured(prisma_edge, gateway_cache) + policies.assert_awaited_once() + policies.return_value = [_POLICY] + 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=None + ): + await invalidate_memory_configuration() + assert await gateway_memory_is_configured(prisma_edge, gateway_cache) + assert policies.await_count == 2 diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_management.py b/tests/test_litellm/proxy/memory/test_memory_v2_management.py new file mode 100644 index 00000000000..b680b4a162e --- /dev/null +++ b/tests/test_litellm/proxy/memory/test_memory_v2_management.py @@ -0,0 +1,254 @@ +"""Exercise real policy administration and self-service logic at the database edge.""" + +from collections.abc import Iterator +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import UI_TEAM_ID, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.memory import management +from litellm.proxy.memory.policy import MemoryIdentity, memory_digest +from litellm.types.memory_v2 import MemoryCapture, MemoryPolicy, MemoryPolicyInput, MemoryPreference + + +@pytest.fixture +def database() -> Iterator[MagicMock]: + client = MagicMock() + for name in ( + "litellm_memorypolicy", + "litellm_memorypreference", + "litellm_memorytable", + "litellm_teamtable", + "litellm_projecttable", + "litellm_organizationtable", + "litellm_organizationmembership", + "litellm_verificationtoken", + "litellm_usertable", + ): + table = getattr(client.db, name) + table.find_unique = AsyncMock(return_value=None) + table.find_first = AsyncMock(return_value=None) + table.find_many = AsyncMock(return_value=[]) + table.upsert = AsyncMock() + table.delete = AsyncMock() + table.delete_many = AsyncMock(return_value=0) + table.create = AsyncMock() + client.db.litellm_teamtable.find_unique.return_value = { + "team_id": "team", + "organization_id": None, + "members_with_roles": [{"user_id": "team-admin", "role": "admin"}], + } + client.db.litellm_projecttable.find_unique.return_value = {"project_id": "project", "team_id": "team"} + client.db.litellm_verificationtoken.find_unique.return_value = { + "token": "a" * 64, + "user_id": "owner", + "team_id": "team", + "org_id": "explicit-org", + } + client.db.litellm_usertable.find_unique.return_value = {"user_id": "owner"} + client.db.litellm_organizationtable.find_unique.return_value = { + "organization_id": "org", + "budget_id": "budget", + "created_by": "admin", + "updated_by": "admin", + } + with patch.multiple( # test-quality-ok: Replace only the external database and cache; exercise actual authorization and repositories. + "litellm.proxy.proxy_server", prisma_client=client, user_api_key_cache=MagicMock(async_delete_cache=AsyncMock()) + ): + yield client + + +def auth(user: str = "owner", role: LitellmUserRoles = LitellmUserRoles.INTERNAL_USER) -> UserAPIKeyAuth: + return UserAPIKeyAuth(token="a" * 64, user_id=user, user_role=role, team_id="team", org_id="explicit-org") + + +def policy(**changes: object) -> MemoryPolicy: + return MemoryPolicy.model_validate( + { + "policy_id": memory_digest("gateway", "*"), + "target_type": "gateway", + "target_id": "*", + "activation": "automatic", + "scope": "key", + "updated_by": "admin", + "updated_at": datetime(2026, 9, 12, tzinfo=timezone.utc), + **changes, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "target,target_id", + [ + ("gateway", "*"), + ("organization", "org"), + ("team", "team"), + ("project", "project"), + ("key", "a" * 64), + ("user", "owner"), + ], +) +async def test_proxy_admin_can_configure_each_existing_target(database: MagicMock, target: str, target_id: str) -> None: + request = MemoryPolicyInput.model_validate( + {"target_type": target, "target_id": target_id, "activation": "opt_in", "scope": "key"} + ) + database.db.litellm_memorypolicy.upsert.return_value = policy(**request.model_dump()) + result = await management.set_policy(request, auth("admin", LitellmUserRoles.PROXY_ADMIN)) + assert result.activation == "opt_in" + written = database.db.litellm_memorypolicy.upsert.call_args.kwargs + assert written["where"]["policy_id"] == memory_digest(target, target_id) + assert written["data"]["create"]["updated_by"] == "admin" + assert written["data"]["update"]["target_id"] == target_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize("target,target_id", [("team", "team"), ("project", "project"), ("key", "a" * 64)]) +async def test_team_admin_can_configure_owned_targets_but_cannot_broaden_to_user_scope( + database: MagicMock, target: str, target_id: str +) -> None: + request = MemoryPolicyInput.model_validate( + {"target_type": target, "target_id": target_id, "activation": "automatic", "scope": "team"} + ) + database.db.litellm_memorypolicy.upsert.return_value = policy(**request.model_dump()) + assert (await management.set_policy(request, auth("team-admin"))).scope == "team" + with pytest.raises(HTTPException) as denied: + await management.set_policy(request.model_copy(update={"scope": "user"}), auth("team-admin")) + assert denied.value.status_code == 403 + database.db.litellm_memorypolicy.upsert.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_org_membership_is_checked_for_the_selected_organization(database: MagicMock) -> None: + request = MemoryPolicyInput( + target_type="organization", target_id="org", activation="automatic", scope="organization" + ) + database.db.litellm_organizationmembership.find_first.return_value = SimpleNamespace(user_role="org_admin") + database.db.litellm_memorypolicy.upsert.return_value = policy(**request.model_dump()) + assert (await management.set_policy(request, auth("org-admin"))).scope == "organization" + assert database.db.litellm_organizationmembership.find_first.call_args.kwargs["where"] == { + "organization_id": "org", + "user_id": "org-admin", + "user_role": "org_admin", + } + database.db.litellm_organizationmembership.find_first.return_value = None + with pytest.raises(HTTPException) as denied: + await management.set_policy(request, auth("org-admin")) + assert denied.value.status_code == 403 + database.db.litellm_memorypolicy.upsert.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "role", + [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY], +) +async def test_non_administrator_cannot_change_gateway_policy(database: MagicMock, role: LitellmUserRoles) -> None: + with pytest.raises(HTTPException) as denied: + await management.set_policy( + MemoryPolicyInput(target_type="gateway", target_id="*", activation="automatic"), auth(role=role) + ) + assert denied.value.status_code == 403 + database.db.litellm_memorypolicy.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_policy_listing_requires_owned_target_and_preserves_pagination(database: MagicMock) -> None: + database.db.litellm_memorypolicy.find_many.return_value = [policy(target_type="team", target_id="team")] + assert len(await management.list_policies("team", "team", 100, auth("team-admin"))) == 1 + assert database.db.litellm_memorypolicy.find_many.call_args.kwargs == { + "where": {"target_type": "team", "target_id": "team"}, + "take": 100, + "skip": 100, + "order": {"policy_id": "asc"}, + } + for caller, target, expected in [(auth(), None, 403), (auth("admin", LitellmUserRoles.PROXY_ADMIN), "team", 400)]: + with pytest.raises(HTTPException) as denied: + await management.list_policies(target, None, 0, caller) + assert denied.value.status_code == expected + assert ( + len(await management.list_policies(None, None, 0, auth("admin", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY))) == 1 + ) + + +@pytest.mark.asyncio +async def test_policy_delete_rechecks_administrator_and_reports_missing(database: MagicMock) -> None: + table = database.db.litellm_memorypolicy + table.find_unique.return_value = policy() + with pytest.raises(HTTPException) as denied: + await management.delete_policy("policy", auth()) + assert denied.value.status_code == 403 + table.delete.assert_not_awaited() + assert (await management.delete_policy("policy", auth("admin", LitellmUserRoles.PROXY_ADMIN))).status_code == 204 + table.find_unique.return_value = None + with pytest.raises(HTTPException) as missing: + await management.delete_policy("policy", auth("admin", LitellmUserRoles.PROXY_ADMIN)) + assert missing.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_preference_updates_are_bound_to_authenticated_subject(database: MagicMock) -> None: + assert not (await management.get_preference(auth())).enabled + table = database.db.litellm_memorypreference + assert (await management.set_preference(MemoryPreference(enabled=True), auth())).enabled + assert table.upsert.call_args.kwargs["where"] == {"subject": memory_digest("user", "owner")} + table.find_unique.return_value = SimpleNamespace(enabled=True) + assert (await management.get_preference(auth())).enabled + assert not (await management.set_preference(MemoryPreference(enabled=False), auth())).enabled + assert table.delete_many.call_args.kwargs["where"] == {"subject": memory_digest("user", "owner")} + with pytest.raises(HTTPException) as denied: + await management.set_preference( + MemoryPreference(enabled=False), auth(role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) + ) + assert denied.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_dashboard_session_can_select_own_key_but_not_another_users_key(database: MagicMock) -> None: + database.db.litellm_memorypolicy.find_many.return_value = [policy()] + session = auth().model_copy(update={"token": "session", "team_id": UI_TEAM_ID}) + own = await management.access_for_key(session, "a" * 64) + assert own.active and own.identity.organization_id == "explicit-org" + assert own.namespace == MemoryIdentity.from_auth(auth()).namespace("key") + assert (await management.get_status(None, auth())).active + database.db.litellm_verificationtoken.find_unique.return_value["user_id"] = "someone-else" + with pytest.raises(HTTPException) as denied: + await management.access_for_key(session, "a" * 64) + assert denied.value.status_code == 403 + database.db.litellm_verificationtoken.find_unique.return_value = None + with pytest.raises(HTTPException) as missing: + await management.access_for_key(auth("admin", LitellmUserRoles.PROXY_ADMIN), "a" * 64) + assert missing.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_entry_endpoints_apply_namespace_and_delete_after_disable(database: MagicMock) -> None: + database.db.litellm_memorypolicy.find_many.return_value = [policy()] + namespace = MemoryIdentity.from_auth(auth()).namespace("key") + table = database.db.litellm_memorytable + assert await management.list_entries("demo", 2, 4, None, auth()) == [] + assert table.find_many.call_args.kwargs["where"]["namespace"] == namespace + now = datetime(2026, 9, 12, tzinfo=timezone.utc) + table.create.return_value = SimpleNamespace( + memory_id="entry", + key=f"memory-v2:{namespace}:demo", + namespace=namespace, + value="Use port 8347", + metadata={"title": "Demo", "evidence": "User said so"}, + updated_at=now, + ) + saved = await management.capture_entry( + MemoryCapture(key="demo", title="Demo", content="Use port 8347", evidence="User said so"), None, auth() + ) + assert saved.content == "Use port 8347" and saved.key == "demo" + database.db.litellm_memorypolicy.find_many.return_value = [policy(activation="disabled")] + table.delete_many.return_value = 1 + assert (await management.delete_entry("entry", None, auth())).status_code == 204 + assert table.delete_many.call_args.kwargs["where"] == {"namespace": namespace, "memory_id": "entry"} + table.delete_many.return_value = 0 + with pytest.raises(HTTPException) as missing: + await management.delete_entry("entry", None, auth()) + assert missing.value.status_code == 404