fix(memory): invalidate worker hints and bound pilot validation

This commit is contained in:
moe-berri 2026-09-12 02:51:40 -07:00
parent 0269841595
commit 7297247aab
6 changed files with 378 additions and 5 deletions

View file

@ -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.

View file

@ -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

View file

@ -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:

View file

@ -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

View file

@ -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

View file

@ -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