diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index 00c4e0070e6..c04bbb92a37 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -71,6 +71,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = ( "/v1/workflows/", "/project/", "/memory/", + "/v2/memory/", "/mcp/", # Control plane (see the List Endpoints + Tables standard). Every resource # eventually moves under this prefix, so allowlist it once rather than diff --git a/deploy/memory-pilot/build.sh b/deploy/memory-pilot/build.sh index 97647e63f37..46daf947233 100755 --- a/deploy/memory-pilot/build.sh +++ b/deploy/memory-pilot/build.sh @@ -4,6 +4,7 @@ set -euo pipefail # writable homes to install the repository's pinned toolchain during uv sync. export RUSTUP_HOME="$PWD/.memory-pilot-rustup" export CARGO_HOME="$PWD/.memory-pilot-cargo" +export PRISMA_BINARY_CACHE_DIR="$PWD/.memory-pilot-prisma" python -m pip install uv==0.11.7 export UV_PROJECT_ENVIRONMENT="$PWD/.memory-pilot-venv" uv sync --frozen --extra proxy --extra extra_proxy --no-default-groups diff --git a/deploy/memory-pilot/pilot.py b/deploy/memory-pilot/pilot.py index 1239c52a932..fc2083fc701 100644 --- a/deploy/memory-pilot/pilot.py +++ b/deploy/memory-pilot/pilot.py @@ -7,15 +7,19 @@ from contextvars import ContextVar from typing import Final import httpx -from fastapi import HTTPException, Request +from starlette.exceptions import HTTPException +from starlette.requests import Request from starlette.responses import JSONResponse from starlette.types import ASGIApp, Receive, Scope, Send from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._types import UI_TEAM_ID, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +from litellm.proxy.auth.user_api_key_auth import _get_bearer_token_or_received_api_key from litellm.repositories.verification_token_repository import VerificationTokenRepository +from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.utils import CallTypesLiteral _UPSTREAM: Final = os.environ["UPSTREAM_LITELLM_BASE_URL"].rstrip("/") @@ -42,25 +46,28 @@ forward_credential: Final = ForwardCredential() class PilotGateway: def __init__(self, app: ASGIApp) -> None: self.app = app - self.upstream = httpx.AsyncClient(timeout=20) + self.upstream = get_async_httpx_client( + httpxSpecialProvider.PassThroughEndpoint, + params={"timeout": 20, "client_alias": "memory-pilot-upstream"}, + ) async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if scope["type"] == "lifespan": try: await self.app(scope, receive, send) finally: - await self.upstream.aclose() + await self.upstream.close() return if scope["type"] != "http": await self.app(scope, receive, send) return request: Final = Request(scope, receive) - credential: Final = ( + credential: Final = _get_bearer_token_or_received_api_key( request.headers.get("x-litellm-api-key") or request.headers.get("authorization") or request.headers.get("x-api-key") or "" - ).removeprefix("Bearer ") + ) from litellm.proxy.proxy_server import master_key, prisma_client if not credential or master_key and secrets.compare_digest(credential, master_key): diff --git a/deploy/memory-pilot/start.sh b/deploy/memory-pilot/start.sh index 37e0af94093..bfce87a3c35 100755 --- a/deploy/memory-pilot/start.sh +++ b/deploy/memory-pilot/start.sh @@ -1,6 +1,8 @@ #!/usr/bin/env bash set -euo pipefail export PATH="$PWD/.memory-pilot-venv/bin:$PATH" +export PRISMA_BINARY_CACHE_DIR="$PWD/.memory-pilot-prisma" +export PRISMA_CLI_PATH="$PRISMA_BINARY_CACHE_DIR/node_modules/.bin/prisma" prisma migrate deploy --schema litellm-proxy-extras/litellm_proxy_extras/schema.prisma export WORKER_CONFIG="$PWD/deploy/memory-pilot/proxy_config.yaml" export PYTHONPATH="$PWD/deploy/memory-pilot${PYTHONPATH:+:$PYTHONPATH}" diff --git a/gateway/routes/allowlist.py b/gateway/routes/allowlist.py index 3733072a948..44583c168f2 100644 --- a/gateway/routes/allowlist.py +++ b/gateway/routes/allowlist.py @@ -74,6 +74,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = ( "/containers", "/v1/evals", "/v1/memory", + "/v2/memory/entries", "/queue/chat/", # Google data plane (v1beta is the Google AI Studio version) "/v1beta/", @@ -123,6 +124,8 @@ GATEWAY_EXACT_PATHS: frozenset[str] = frozenset( "/redoc", "/test", "/debug/memory/summary", + "/v2/memory/status", + "/v2/memory/preference", } ) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260912000000_gateway_memory_v2/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260912000000_gateway_memory_v2/migration.sql index 160e4cc38f8..f6ba9da0d2c 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260912000000_gateway_memory_v2/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260912000000_gateway_memory_v2/migration.sql @@ -1,9 +1,9 @@ -ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN "namespace" TEXT; +ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN IF NOT EXISTS "namespace" TEXT; -CREATE INDEX "LiteLLM_MemoryTable_namespace_updated_at_idx" +CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_namespace_updated_at_idx" ON "LiteLLM_MemoryTable"("namespace", "updated_at"); -CREATE TABLE "LiteLLM_MemoryPolicy" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_MemoryPolicy" ( "policy_id" TEXT NOT NULL, "target_type" TEXT NOT NULL, "target_id" TEXT NOT NULL, @@ -15,10 +15,10 @@ CREATE TABLE "LiteLLM_MemoryPolicy" ( CONSTRAINT "LiteLLM_MemoryPolicy_pkey" PRIMARY KEY ("policy_id") ); -CREATE UNIQUE INDEX "LiteLLM_MemoryPolicy_target_type_target_id_key" +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MemoryPolicy_target_type_target_id_key" ON "LiteLLM_MemoryPolicy"("target_type", "target_id"); -CREATE TABLE "LiteLLM_MemoryPreference" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_MemoryPreference" ( "subject" TEXT NOT NULL, "enabled" BOOLEAN NOT NULL DEFAULT false, "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c398abff099..5df0fbf8d73 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -522,6 +522,7 @@ class RequestRateLimiterStash: owner_litellm_call_id: str | None = None rate_limit_response: RateLimitResponse | None = None parallel_slot: ParallelSlotAcquisition | None = None + parallel_release_complete: asyncio.Event = field(default_factory=asyncio.Event) reserved_tokens: int = 0 reserved_model: str | None = None reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) @@ -548,6 +549,13 @@ def get_request_stash() -> RequestRateLimiterStash | None: return _request_stash.get() +async def wait_for_request_parallel_release() -> None: + """Let sequential internal requests wait for their deferred slot release.""" + stash: Final = get_request_stash() + if stash is not None and stash.parallel_slot is not None: + await stash.parallel_release_complete.wait() + + def get_or_create_request_stash() -> RequestRateLimiterStash: stash = _request_stash.get() if stash is None: @@ -3375,6 +3383,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) stash.parallel_slot = None + stash.parallel_release_complete.set() self._handle_rate_limit_error( response=io_response, descriptors=descriptors, @@ -3547,6 +3556,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): slot_id=parallel_slot_id, counter_keys=parallel_counter_keys, ) + stash.parallel_release_complete.clear() # ---------------------------------------------------------------- # TPM token reservation @@ -3638,6 +3648,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) stash.parallel_slot = None + stash.parallel_release_complete.set() self._handle_rate_limit_error( response=tpm_response, descriptors=descriptors, @@ -4457,6 +4468,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=litellm_parent_otel_span, ) stash.parallel_slot = None + stash.parallel_release_complete.set() pipeline_operations: Final = self._build_success_event_pipeline_operations( kwargs=kwargs, @@ -4583,6 +4595,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=litellm_parent_otel_span, ) stash.parallel_slot = None + stash.parallel_release_complete.set() # Skip the reservation refund if async_post_call_failure_hook # already released it (proxy-level rejection that also bubbles up @@ -4699,6 +4712,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=None, ) stash.parallel_slot = None + stash.parallel_release_complete.set() async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response): """ @@ -4780,6 +4794,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) stash.parallel_slot = None + stash.parallel_release_complete.set() if stash.batch_enqueued_reservation is not None: await self.batch_enqueued_token_store.refund( diff --git a/litellm/proxy/memory/gateway.py b/litellm/proxy/memory/gateway.py index c2a9889baee..221a0bc5b5f 100644 --- a/litellm/proxy/memory/gateway.py +++ b/litellm/proxy/memory/gateway.py @@ -20,7 +20,8 @@ from litellm.litellm_core_utils.prompt_templates.server_tools import ( prepare_server_tools, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.memory.policy import MemoryIdentity, resolve_memory_access +from litellm.proxy.hooks.parallel_request_limiter_v3 import wait_for_request_parallel_release +from litellm.proxy.memory.policy import MemoryIdentity, gateway_memory_is_configured, resolve_memory_access from litellm.proxy.memory.store import MemoryStore from litellm.types.memory_v2 import MemoryCapture, MemoryRead, MemorySearch @@ -155,11 +156,16 @@ async def prepare_gateway_memory( ) -> dict[str, object]: if _memory_call.get() or route not in ("acompletion", "aresponses", "anthropic_messages"): return data - from litellm.proxy.proxy_server import app, prisma_client + from litellm.proxy.proxy_server import app, prisma_client, user_api_key_cache if prisma_client is None: return data - access: Final = await resolve_memory_access(prisma_client, MemoryIdentity.from_auth(auth)) + identity: Final = MemoryIdentity.from_auth(auth) + if not identity.user_id and not identity.key_id: + return data + if not await gateway_memory_is_configured(prisma_client, user_api_key_cache): + return data + access: Final = await resolve_memory_access(prisma_client, identity) if not access.active: return data functions: Final = tuple( @@ -184,18 +190,35 @@ async def prepare_gateway_memory( } token: Final = _memory_call.set(True) try: - async with httpx.AsyncClient( - transport=httpx.ASGITransport(app=app), base_url="http://litellm-memory" - ) as client: + # Dispatch in process through the existing authenticated endpoint. No + # network client, TLS context, or connection pool is created here. + async with httpx.ASGITransport(app=app) as transport: + + async def dispatch_round(body: Mapping[str, object]) -> httpx.Response: + result: Final = await transport.handle_async_request( + httpx.Request( + "POST", + "http://litellm-memory" + request.url.path, + json=body, + headers=headers, + params=request.query_params, + ) + ) + await result.aread() + # Success accounting runs asynchronously. The next model round + # must not compete with this completed call for the same slot. + await wait_for_request_parallel_release() + return result async def call_model(body: Mapping[str, object]) -> Mapping[str, object]: round_body: Final = { # mutable-ok: HTTP JSON serialization requires a native dictionary. **body, "litellm_call_id": str(uuid4()), } - result: Final = await client.post( - request.url.path, json=round_body, headers=headers, params=request.query_params - ) + # Each endpoint owns its request context, including the rate + # limiter's mutable stash. Reusing this task would let the next + # round overwrite the owner seen by deferred logging callbacks. + result: Final = await asyncio.create_task(dispatch_round(round_body)) if result.is_error: raise HTTPException( status_code=result.status_code, diff --git a/litellm/proxy/memory/management.py b/litellm/proxy/memory/management.py index 65ac4b98c10..40557981bf6 100644 --- a/litellm/proxy/memory/management.py +++ b/litellm/proxy/memory/management.py @@ -5,7 +5,14 @@ from fastapi import APIRouter, Depends, HTTPException, Query, Response from litellm.proxy._types import UI_TEAM_ID, LitellmUserRoles, UserAPIKeyAuth, user_api_key_has_admin_view from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.memory.memory_endpoints import is_memory_team_admin, require_memory_prisma -from litellm.proxy.memory.policy import MemoryAccess, MemoryIdentity, memory_digest, resolve_memory_access +from litellm.proxy.memory.policy import ( + MemoryAccess, + MemoryIdentity, + invalidate_memory_configuration, + memory_digest, + memory_primary_client, + resolve_memory_access, +) from litellm.proxy.memory.store import MemoryStore from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.project_repository import ProjectRepository @@ -41,7 +48,7 @@ router: Final = APIRouter( async def require_policy_admin( auth: UserAPIKeyAuth, target_type: MemoryTarget, target_id: str, *, write: bool = True ) -> None: - prisma: Final = require_memory_prisma() + prisma: Final = memory_primary_client(require_memory_prisma()) if write and auth.user_role in (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY): raise HTTPException(status_code=403, detail="Memory policies require administrator write access") proxy_admin: Final = ( @@ -92,7 +99,7 @@ async def list_policies( raise HTTPException(status_code=403, detail="Select a target you administer") elif target_type is not None or target_id is not None: raise HTTPException(status_code=400, detail="Provide both target_type and target_id") - rows: Final = await MemoryPolicyRepository(require_memory_prisma()).table.find_many( + rows: Final = await MemoryPolicyRepository(memory_primary_client(require_memory_prisma())).table.find_many( where={ # mutable-ok: Prisma serializes these as native JSON containers. "target_type": target_type, "target_id": target_id, @@ -121,7 +128,7 @@ async def set_policy(policy: MemoryPolicyInput, auth: UserAPIKeyAuth = _AUTH) -> **policy.model_dump(), "updated_by": auth.user_id or "proxy-admin", } - row: Final = await MemoryPolicyRepository(require_memory_prisma()).table.upsert( + row: Final = await MemoryPolicyRepository(memory_primary_client(require_memory_prisma())).table.upsert( where={ # mutable-ok: Prisma serializes these as native JSON containers. "policy_id": policy_id }, @@ -133,12 +140,13 @@ async def set_policy(policy: MemoryPolicyInput, auth: UserAPIKeyAuth = _AUTH) -> "update": fields, }, ) + await invalidate_memory_configuration() return MemoryPolicy.model_validate(row, from_attributes=True) @router.delete("/policies/{policy_id}", status_code=204) async def delete_policy(policy_id: str, auth: UserAPIKeyAuth = _AUTH) -> Response: - table: Final = MemoryPolicyRepository(require_memory_prisma()).table + table: Final = MemoryPolicyRepository(memory_primary_client(require_memory_prisma())).table row: Final = await table.find_unique( where={ # mutable-ok: Prisma serializes these as native JSON containers. "policy_id": policy_id @@ -153,13 +161,14 @@ async def delete_policy(policy_id: str, auth: UserAPIKeyAuth = _AUTH) -> Respons "policy_id": policy_id } ) + await invalidate_memory_configuration() return Response(status_code=204) @router.get("/preference", response_model=MemoryPreference) async def get_preference(auth: UserAPIKeyAuth = _AUTH) -> MemoryPreference: subject: Final = MemoryIdentity.from_auth(auth).preference_subject - row: Final = await MemoryPreferenceRepository(require_memory_prisma()).table.find_unique( + row: Final = await MemoryPreferenceRepository(memory_primary_client(require_memory_prisma())).table.find_unique( where={ # mutable-ok: Prisma serializes these as native JSON containers. "subject": subject } @@ -174,13 +183,13 @@ async def set_preference(preference: MemoryPreference, auth: UserAPIKeyAuth = _A raise HTTPException(status_code=403, detail="Read-only users cannot change memory preferences") subject: Final = identity.preference_subject if not preference.enabled: - await MemoryPreferenceRepository(require_memory_prisma()).table.delete_many( + await MemoryPreferenceRepository(memory_primary_client(require_memory_prisma())).table.delete_many( where={ # mutable-ok: Prisma serializes these as native JSON containers. "subject": subject } ) return preference - await MemoryPreferenceRepository(require_memory_prisma()).table.upsert( + await MemoryPreferenceRepository(memory_primary_client(require_memory_prisma())).table.upsert( where={ # mutable-ok: Prisma serializes these as native JSON containers. "subject": subject }, @@ -206,7 +215,7 @@ async def get_status( async def access_for_key(auth: UserAPIKeyAuth, key_id: str | None) -> MemoryAccess: - prisma: Final = require_memory_prisma() + prisma: Final = memory_primary_client(require_memory_prisma()) if key_id is None: return await resolve_memory_access(prisma, MemoryIdentity.from_auth(auth)) key: Final = await VerificationTokenRepository(prisma).find_by_id(key_id, id_field="token") @@ -224,7 +233,7 @@ async def access_for_key(auth: UserAPIKeyAuth, key_id: str | None) -> MemoryAcce user_id=key.user_id, team_id=key.team_id, project_id=key.project_id, - organization_id=team.organization_id if team else key.org_id, + organization_id=key.org_id or (team.organization_id if team else None), read_only=MemoryIdentity.from_auth(auth).read_only, ) return await resolve_memory_access(prisma, identity) @@ -238,7 +247,7 @@ async def list_entries( key_id: str | None = Query(None, pattern=r"^[a-f0-9]{64}$"), auth: UserAPIKeyAuth = _AUTH, ) -> list[MemoryEntry]: - prisma: Final = require_memory_prisma() + prisma: Final = memory_primary_client(require_memory_prisma()) access: Final = await access_for_key(auth, key_id) return await MemoryStore(prisma, access).search( MemorySearch(query=query, limit=limit, offset=offset), require_active=False @@ -251,7 +260,7 @@ async def capture_entry( key_id: str | None = Query(None, pattern=r"^[a-f0-9]{64}$"), auth: UserAPIKeyAuth = _AUTH, ) -> MemoryEntry: - prisma: Final = require_memory_prisma() + prisma: Final = memory_primary_client(require_memory_prisma()) access: Final = await access_for_key(auth, key_id) return await MemoryStore(prisma, access).capture(capture) @@ -262,7 +271,7 @@ async def delete_entry( key_id: str | None = Query(None, pattern=r"^[a-f0-9]{64}$"), auth: UserAPIKeyAuth = _AUTH, ) -> Response: - prisma: Final = require_memory_prisma() + prisma: Final = memory_primary_client(require_memory_prisma()) access: Final = await access_for_key(auth, key_id) if not await MemoryStore(prisma, access).delete(memory_id): raise HTTPException(status_code=404, detail="Memory not found") diff --git a/litellm/proxy/memory/policy.py b/litellm/proxy/memory/policy.py index 7fd53f01333..6ca18cc4909 100644 --- a/litellm/proxy/memory/policy.py +++ b/litellm/proxy/memory/policy.py @@ -5,10 +5,46 @@ from typing import Final from fastapi import HTTPException +from litellm.caching.caching import DualCache from litellm.proxy._types import UI_TEAM_ID, LitellmUserRoles, UserAPIKeyAuth +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 +_CONFIGURED_CACHE_KEY: Final = "litellm:memory_v2:configured" + + +async def gateway_memory_is_configured(prisma_client: object, cache: DualCache) -> bool: + """Avoid database work on ordinary requests when memory is not configured. + + This is only a presence hint. Authorization is always checked against the + primary before memory is used. The short TTL also covers separate workers + 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 + 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) + return configured + + +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) + + +def memory_primary_client(prisma_client: object) -> WriterPinnedClient: + """Memory authorization and read-after-write must never use a lagging replica.""" + db: Final = getattr(prisma_client, "db", None) + if db is None: + raise RuntimeError("Memory requires a connected Prisma database") + # Pin the actual writer even while unavailable: memory must fail closed + # instead of authorizing storage or recall from stale policy rows. + return WriterPinnedClient(db.writer if isinstance(db, RoutingPrismaWrapper) else db) + def memory_digest(*parts: str | None) -> str: return hashlib.sha256(json.dumps(parts, separators=(",", ":")).encode()).hexdigest() @@ -119,7 +155,8 @@ class MemoryAccess: async def resolve_memory_access(prisma_client: object, identity: MemoryIdentity) -> MemoryAccess: - rows: Final = await MemoryPolicyRepository(prisma_client).table.find_many( + primary: Final = memory_primary_client(prisma_client) + rows: Final = await MemoryPolicyRepository(primary).table.find_many( where={ # mutable-ok: Prisma serializes these as native JSON containers. "OR": [ # mutable-ok: Prisma serializes these as native JSON containers. { # mutable-ok: Prisma serializes these as native JSON containers. @@ -137,7 +174,7 @@ async def resolve_memory_access(prisma_client: object, identity: MemoryIdentity) policy: Final = next((policies[target] for target in identity.policy_targets if target in policies), None) if not identity.user_id and not identity.key_id: return MemoryAccess(identity=identity, policy=None, opted_in=False) - preference: Final = await MemoryPreferenceRepository(prisma_client).table.find_unique( + preference: Final = await MemoryPreferenceRepository(primary).table.find_unique( where={ # mutable-ok: Prisma serializes these as native JSON containers. "subject": identity.preference_subject } diff --git a/litellm/proxy/memory/store.py b/litellm/proxy/memory/store.py index 42201cc3cb7..db3e61bd174 100644 --- a/litellm/proxy/memory/store.py +++ b/litellm/proxy/memory/store.py @@ -2,10 +2,9 @@ import json from typing import TYPE_CHECKING, Final from fastapi import HTTPException -from prisma.errors import UniqueViolationError from pydantic import TypeAdapter -from litellm.proxy.memory.policy import MemoryAccess, memory_digest, resolve_memory_access +from litellm.proxy.memory.policy import MemoryAccess, memory_digest, memory_primary_client, resolve_memory_access from litellm.repositories.table_repositories import MemoryRepository from litellm.types.memory_v2 import MemoryCapture, MemoryEntry, MemorySearch @@ -36,9 +35,9 @@ def memory_entry(row: "LiteLLM_MemoryTable") -> MemoryEntry: class MemoryStore: def __init__(self, prisma_client: object, access: MemoryAccess) -> None: - self.prisma_client = prisma_client + self.prisma_client = memory_primary_client(prisma_client) self.access = access - self.table = MemoryRepository(prisma_client).table + self.table = MemoryRepository(self.prisma_client).table async def _namespace(self, *, write: bool = False, require_active: bool = True) -> str: current: Final = await resolve_memory_access(self.prisma_client, self.access.identity) @@ -115,6 +114,8 @@ class MemoryStore: return memory_entry(row) async def capture(self, capture: MemoryCapture) -> MemoryEntry: + from prisma.errors import UniqueViolationError + namespace: Final = await self._namespace(write=True) key: Final = f"memory-v2:{namespace}:{capture.key}" metadata: Final = { # mutable-ok: Prisma serializes these as native JSON containers. diff --git a/litellm/types/memory_v2.py b/litellm/types/memory_v2.py index 8a8597c6d64..2db5bfd0056 100644 --- a/litellm/types/memory_v2.py +++ b/litellm/types/memory_v2.py @@ -1,7 +1,8 @@ from datetime import datetime -from typing import Literal, Self, TypeAlias +from typing import Literal, TypeAlias from pydantic import BaseModel, ConfigDict, Field, model_validator +from typing_extensions import Self MemoryTarget: TypeAlias = Literal["gateway", "organization", "team", "project", "user", "key"] MemoryScope: TypeAlias = Literal["key", "user", "team", "project", "organization"] diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index e8caa801467..37e758e149f 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -306,6 +306,7 @@ CONTROL_PLANE_PREFIXES: tuple[str, ...] = ( "/global", "/config", "/guardrails", + "/v2/memory/policies", "/openapi.json", ) diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/proxy_unit_tests/test_jwt.py index 6ad253f33e8..557bb4e9b1f 100644 --- a/tests/proxy_unit_tests/test_jwt.py +++ b/tests/proxy_unit_tests/test_jwt.py @@ -1182,7 +1182,11 @@ async def test_end_user_jwt_auth(monkeypatch): "llm_router", router, ) - setattr(litellm.proxy.proxy_server, "prisma_client", {}) + # No memory policy exists for this authenticated JWT user. + memory_db = MagicMock() + memory_db.db.litellm_memorypolicy.find_many = AsyncMock(return_value=[]) + memory_db.db.litellm_memorypreference.find_unique = AsyncMock(return_value=None) + setattr(litellm.proxy.proxy_server, "prisma_client", memory_db) setattr(litellm.proxy.proxy_server, "jwt_handler", jwt_handler) from litellm.proxy.proxy_server import cost_tracking diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index d06eb0426c9..ee8b45694ce 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -503,7 +503,11 @@ def test_custom_logger_failure_handler(mock_acompletion, client_no_auth): setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") - setattr(litellm.proxy.proxy_server, "prisma_client", "FAKE-VAR") + # The request has no memory policy; keep the database edge inert. + memory_db = MagicMock() + memory_db.db.litellm_memorypolicy.find_many = AsyncMock(return_value=[]) + memory_db.db.litellm_memorypreference.find_unique = AsyncMock(return_value=None) + setattr(litellm.proxy.proxy_server, "prisma_client", memory_db) setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging_obj) with patch.object( diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 10c0bb88a82..1961e67e420 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4232,6 +4232,12 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch): await local_cache.async_get_cache(key=counter_key) ) == 1 + from litellm.proxy.hooks.parallel_request_limiter_v3 import wait_for_request_parallel_release + + waiter = asyncio.create_task(wait_for_request_parallel_release()) + await asyncio.sleep(0) + assert not waiter.done(), "Internal rounds must wait while the completed call still owns a slot" + await handler.async_log_success_event( kwargs={ "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, @@ -4242,6 +4248,7 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch): start_time=datetime.now(), end_time=datetime.now(), ) + await asyncio.wait_for(waiter, timeout=1) assert handler._gauge_in_flight_from_cache_value( await local_cache.async_get_cache(key=counter_key) ) == 0 diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py new file mode 100644 index 00000000000..4e2481e2c86 --- /dev/null +++ b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py @@ -0,0 +1,366 @@ +"""Failure and authorization boundaries, with only the database/model edges replaced.""" + +import json +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException +from prisma.models import LiteLLM_MemoryTable + +from litellm.proxy.memory.gateway import execute_memory_tool, run_memory_tools +from litellm.proxy.memory.policy import MemoryAccess, MemoryIdentity, resolve_memory_access +from litellm.proxy.memory.store import MemoryStore +from litellm.types.memory_v2 import MemoryCapture, MemoryPolicy, MemorySearch + +_NOW: Final = datetime(2026, 9, 12, tzinfo=timezone.utc) +_IDENTITY: Final = MemoryIdentity("a" * 64, "owner", "team", "project", "org", False) +_POLICY: Final = MemoryPolicy( + policy_id="policy", + target_type="team", + target_id="team", + activation="automatic", + scope="key", + updated_at=_NOW, + updated_by="admin", +) +_CAPTURE: Final = MemoryCapture(key="demo", title="Demo", content="Use port 8347", evidence="User selected this port") + + +@pytest.fixture +def prisma_edge() -> MagicMock: + client = MagicMock() + client.db.litellm_memorypolicy.find_many = AsyncMock(return_value=[_POLICY]) + client.db.litellm_memorypreference.find_unique = AsyncMock(return_value=None) + table = client.db.litellm_memorytable + table.find_unique = AsyncMock(return_value=None) + table.find_first = AsyncMock(return_value=None) + table.find_many = AsyncMock(return_value=[]) + table.create = AsyncMock() + table.update_many = AsyncMock(return_value=1) + table.delete_many = AsyncMock(return_value=1) + return client + + +def store(client: MagicMock, identity: MemoryIdentity = _IDENTITY) -> MemoryStore: + return MemoryStore(client, MemoryAccess(identity, _POLICY, False)) + + +def row(**changes: object) -> LiteLLM_MemoryTable: + return LiteLLM_MemoryTable.model_validate( + { + "memory_id": "entry", + "key": f"memory-v2:{_IDENTITY.namespace('key')}:demo", + "namespace": _IDENTITY.namespace("key"), + "value": _CAPTURE.content, + "metadata": json.dumps({"title": _CAPTURE.title, "evidence": _CAPTURE.evidence}), + "created_at": _NOW, + "updated_at": _NOW, + **changes, + } + ) + + +@pytest.mark.asyncio +async def test_policy_precedence_and_opt_in_are_resolved_from_database(prisma_edge: MagicMock) -> None: + team = _POLICY.model_copy(update={"activation": "opt_in"}) + key = _POLICY.model_copy(update={"target_type": "key", "target_id": "a" * 64, "activation": "disabled"}) + policies = prisma_edge.db.litellm_memorypolicy.find_many + policies.return_value = [team, key] + prisma_edge.db.litellm_memorypreference.find_unique.return_value = SimpleNamespace(enabled=True) + access = await resolve_memory_access(prisma_edge, _IDENTITY) + assert access.policy == key and not access.active and access.opted_in + policies.return_value = [team] + assert (await resolve_memory_access(prisma_edge, _IDENTITY)).active + prisma_edge.db.litellm_memorypreference.find_unique.return_value = None + assert not (await resolve_memory_access(prisma_edge, _IDENTITY)).active + policies.return_value = [] + assert (await resolve_memory_access(prisma_edge, _IDENTITY)).namespace is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["disabled", "scope", "missing", "readonly"]) +async def test_store_rechecks_policy_before_writing(prisma_edge: MagicMock, change: str) -> None: + identity = _IDENTITY + if change == "disabled": + prisma_edge.db.litellm_memorypolicy.find_many.return_value = [ + _POLICY.model_copy(update={"activation": "disabled"}) + ] + elif change == "scope": + prisma_edge.db.litellm_memorypolicy.find_many.return_value = [_POLICY.model_copy(update={"scope": "team"})] + elif change == "missing": + prisma_edge.db.litellm_memorypolicy.find_many.return_value = [] + else: + identity = MemoryIdentity("a" * 64, "owner", "team", "project", "org", True) + with pytest.raises(HTTPException) as exc: + await store(prisma_edge, identity).capture(_CAPTURE) + assert exc.value.status_code == 403 + prisma_edge.db.litellm_memorytable.create.assert_not_awaited() + prisma_edge.db.litellm_memorytable.update_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_search_requires_namespace_and_all_words_with_bounded_paging(prisma_edge: MagicMock) -> None: + prisma_edge.db.litellm_memorytable.find_many.return_value = [row()] + entries = await store(prisma_edge).search(MemorySearch(query="port demo port", limit=2, offset=3)) + assert entries[0].content == "Use port 8347" + query = prisma_edge.db.litellm_memorytable.find_many.call_args.kwargs + assert query["where"]["namespace"] == _IDENTITY.namespace("key") + assert query["take"] == 2 and query["skip"] == 3 + clauses = query["where"]["AND"] + assert len(clauses) == 2 + assert [clause["OR"][0]["value"]["contains"] for clause in clauses] == ["port", "demo"] + assert query["order"] == [{"updated_at": "desc"}, {"memory_id": "asc"}] + + +@pytest.mark.asyncio +async def test_read_and_delete_cannot_address_another_namespace(prisma_edge: MagicMock) -> None: + memory = store(prisma_edge) + with pytest.raises(HTTPException) as exc: + await memory.read("foreign-entry") + assert exc.value.status_code == 404 + assert prisma_edge.db.litellm_memorytable.find_first.call_args.kwargs["where"] == { + "namespace": _IDENTITY.namespace("key"), + "memory_id": "foreign-entry", + } + prisma_edge.db.litellm_memorypolicy.find_many.return_value = [_POLICY.model_copy(update={"activation": "disabled"})] + assert await memory.delete("entry") + assert prisma_edge.db.litellm_memorytable.delete_many.call_args.kwargs["where"] == { + "namespace": _IDENTITY.namespace("key"), + "memory_id": "entry", + } + with pytest.raises(HTTPException) as inactive: + await memory.read("entry") + assert inactive.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_identical_capture_is_idempotent_and_new_capture_has_scoped_identity(prisma_edge: MagicMock) -> None: + table = prisma_edge.db.litellm_memorytable + table.create.return_value = row() + saved = await store(prisma_edge).capture(_CAPTURE) + assert saved.content == _CAPTURE.content + data = table.create.call_args.kwargs["data"] + assert data["namespace"] == _IDENTITY.namespace("key") and data["user_id"] == "owner" and data["team_id"] == "team" + assert data["key"].startswith("memory-v2:" + _IDENTITY.namespace("key") + ":") + table.find_unique.return_value = row() + assert await store(prisma_edge).capture(_CAPTURE) == saved + table.create.assert_awaited_once() + table.update_many.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("race", ["missing", "foreign", "stale", "concurrent"]) +async def test_capture_rejects_stale_or_conflicting_replacements(prisma_edge: MagicMock, race: str) -> None: + table = prisma_edge.db.litellm_memorytable + table.find_unique.return_value = ( + None if race == "missing" else row(namespace="foreign" if race == "foreign" else _IDENTITY.namespace("key")) + ) + table.update_many.return_value = 0 if race == "concurrent" else 1 + correction = _CAPTURE.model_copy( + update={ + "content": "Use port 8348", + "expected_revision": _NOW - timedelta(seconds=1) if race == "stale" else _NOW, + } + ) + with pytest.raises(HTTPException) as exc: + await store(prisma_edge).capture(correction) + assert exc.value.status_code == 409 + table.create.assert_not_awaited() + if race == "concurrent": + where = table.update_many.call_args.kwargs["where"] + assert ( + where["namespace"] == _IDENTITY.namespace("key") + and where["updated_at"] == _NOW + and where["value"] == _CAPTURE.content + ) + assert "metadata" in where + else: + table.update_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_successful_replacement_reads_confirmed_updated_row(prisma_edge: MagicMock) -> None: + table = prisma_edge.db.litellm_memorytable + table.find_unique.return_value = row() + table.find_first.return_value = row(value="Use port 8348", updated_at=_NOW + timedelta(seconds=1)) + result = await store(prisma_edge).capture( + _CAPTURE.model_copy(update={"content": "Use port 8348", "expected_revision": _NOW}) + ) + assert result.content == "Use port 8348" and result.updated_at > _NOW + table.update_many.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_tool_argument_errors_are_recoverable_but_revocation_aborts(prisma_edge: MagicMock) -> None: + memory = store(prisma_edge) + invalid = await execute_memory_tool( + memory, {"id": "a", "name": "litellm_memory_capture", "arguments": {"key": "missing-fields"}} + ) + assert invalid.context == "" and "error" in invalid.output + missing = await execute_memory_tool( + memory, {"id": "a", "name": "litellm_memory_read", "arguments": {"memory_id": "missing"}} + ) + assert missing.output == {"error": "Memory not found", "status": 404} + unknown = await execute_memory_tool(memory, {"id": "a", "name": "other_tool", "arguments": {}}) + assert unknown.output == {"error": "Unknown memory tool"} + prisma_edge.db.litellm_memorypolicy.find_many.return_value = [] + with pytest.raises(HTTPException) as exc: + await execute_memory_tool(memory, {"id": "a", "name": "litellm_memory_search", "arguments": {}}) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_model_loop_is_bounded_and_search_results_reach_followup(prisma_edge: MagicMock) -> None: + prisma_edge.db.litellm_memorytable.find_many.return_value = [row()] + model = AsyncMock( + return_value={"content": [{"type": "tool_use", "id": "search", "name": "litellm_memory_search", "input": {}}]} + ) + context = await run_memory_tools( + {"messages": [{"role": "user", "content": "My demo port?"}]}, "anthropic_messages", store(prisma_edge), model + ) + assert model.await_count == 3 and len(context) == 3 + assert all("8347" in reference for reference in context) + continuation = model.call_args_list[1].args[0]["messages"] + assert continuation[-1]["content"][0]["tool_use_id"] == "search" + assert "8347" in continuation[-1]["content"][0]["content"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_id,count", [(True, 1), (False, 9)]) +async def test_invalid_model_calls_are_rejected_before_storage( + prisma_edge: MagicMock, bad_id: bool, count: int +) -> None: + model = AsyncMock( + return_value={ + "content": [ + {"type": "tool_use", "id": "" if bad_id else str(i), "name": "litellm_memory_search", "input": {}} + for i in range(count) + ] + } + ) + with pytest.raises(HTTPException) as exc: + await run_memory_tools({"messages": []}, "anthropic_messages", store(prisma_edge), model) + assert exc.value.status_code == 502 + prisma_edge.db.litellm_memorytable.find_many.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("writer_unavailable", [False, True]) +async def test_replica_lag_cannot_authorize_memory_after_primary_revocation( + prisma_edge: MagicMock, writer_unavailable: bool +) -> None: + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + writer = MagicMock(spec=PrismaWrapper) + reader = MagicMock(spec=PrismaWrapper) + writer.litellm_memorypolicy = SimpleNamespace( + find_many=AsyncMock(return_value=[_POLICY.model_copy(update={"activation": "disabled"})]) + ) + reader.litellm_memorypolicy = SimpleNamespace(find_many=AsyncMock(return_value=[_POLICY])) + writer.litellm_memorypreference = SimpleNamespace(find_unique=AsyncMock(return_value=None)) + writer.litellm_memorytable = prisma_edge.db.litellm_memorytable + writer.is_connected = MagicMock(return_value=False) + reader.is_connected = MagicMock(return_value=False) + routed = RoutingPrismaWrapper(writer, reader) + if writer_unavailable: + writer.connect = AsyncMock(side_effect=RuntimeError("primary unavailable")) + reader.connect = AsyncMock() + await routed.connect() + client = SimpleNamespace(db=routed) + access = await resolve_memory_access(client, _IDENTITY) + assert not access.active + reader.litellm_memorypolicy.find_many.assert_not_awaited() + with pytest.raises(HTTPException) as exc: + await MemoryStore(client, MemoryAccess(_IDENTITY, _POLICY, False)).capture(_CAPTURE) + assert exc.value.status_code == 403 + writer.litellm_memorytable.create.assert_not_awaited() + + +@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 + + cache = DualCache() + policies = prisma_edge.db.litellm_memorypolicy.find_many + policies.return_value = [] + assert not await gateway_memory_is_configured(prisma_edge, cache) + assert not await gateway_memory_is_configured(prisma_edge, cache) + policies.assert_awaited_once() + assert policies.call_args.kwargs == {"take": 1} + enabled_cache = DualCache() + policies.return_value = [_POLICY] + assert await gateway_memory_is_configured(prisma_edge, enabled_cache) + policies.return_value = [_POLICY.model_copy(update={"activation": "disabled"})] + assert await gateway_memory_is_configured(prisma_edge, enabled_cache) + assert not (await resolve_memory_access(prisma_edge, _IDENTITY)).active + + +@pytest.mark.asyncio +async def test_gateway_rounds_keep_separate_limiter_contexts_and_original_client_tools(prisma_edge: MagicMock) -> None: + import asyncio + from unittest.mock import patch + + from fastapi import FastAPI, Request + + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.parallel_request_limiter_v3 import claim_request_stash_for_data, get_request_stash + from litellm.proxy.memory.gateway import prepare_gateway_memory + + provider = FastAPI() + observed = [] + deferred = [] + release = asyncio.Event() + prisma_edge.db.litellm_memorytable.find_many.return_value = [row()] + + @provider.post("/v1/messages") + async def model(request: Request): + body = await request.json() + call_id = body["litellm_call_id"] + stash = claim_request_stash_for_data(body) + observed.append((call_id, stash, body)) + + async def logged_owner(): + await release.wait() + return get_request_stash().owner_litellm_call_id + + deferred.append(asyncio.create_task(logged_owner())) + return {"content": [{"type": "tool_use", "id": "search", "name": "litellm_memory_search", "input": {}}]} + + original = { + "model": "demo", + "stream": True, + "max_tokens": 1, + "messages": [{"role": "user", "content": "My port?"}], + "tools": [{"name": "client_tool"}], + "tool_choice": {"type": "tool", "name": "client_tool"}, + } + request = Request( + { + "type": "http", + "path": "/v1/messages", + "query_string": b"", + "headers": [(b"authorization", b"Bearer test"), (b"idempotency-key", b"visible-only")], + } + ) + 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()): + prepared = await prepare_gateway_memory(original, request, auth, "anthropic_messages") + release.set() + owners = await asyncio.gather(*deferred) + assert len(observed) == 3 and len({id(stash) for _, stash, _ in observed}) == 3 + assert owners == [call_id for call_id, _, _ in observed] + assert len(set(owners)) == 3 + assert get_request_stash() is before + assert all(body["stream"] is False and body["max_tokens"] == 2048 for _, _, body in observed) + 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) diff --git a/tests/test_litellm/proxy/test_component_allowlists.py b/tests/test_litellm/proxy/test_component_allowlists.py index 3073908fa54..3faaf194ee8 100644 --- a/tests/test_litellm/proxy/test_component_allowlists.py +++ b/tests/test_litellm/proxy/test_component_allowlists.py @@ -235,3 +235,14 @@ def test_every_app_mount_is_assigned_to_a_component(): f"Add them to GATEWAY_MOUNT_PATHS, BACKEND_MOUNT_PATHS, or serve them " f"from the UI container:\n " + "\n ".join(sorted(unassigned)) ) + + +def test_memory_v2_policies_stay_on_backend_while_own_entries_are_available_on_gateway(): + gateway = _component_paths(app.router.routes, GATEWAY_EXACT_PATHS, GATEWAY_PATH_PREFIXES) + backend = _component_paths(app.router.routes, BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES) + for path in ("/v2/memory/policies", "/v2/memory/policies/{policy_id}"): + assert path in backend + assert path not in gateway + for path in ("/v2/memory/status", "/v2/memory/preference", "/v2/memory/entries", "/v2/memory/entries/{memory_id}"): + assert path in gateway + assert path in backend diff --git a/ui/litellm-dashboard/src/components/leftnav.test.tsx b/ui/litellm-dashboard/src/components/leftnav.test.tsx index 6eb0218c41d..66400c6c58f 100644 --- a/ui/litellm-dashboard/src/components/leftnav.test.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.test.tsx @@ -384,7 +384,7 @@ describe("Sidebar (leftnav)", () => { }); }); - // Workflow Runs, Memory and Guardrails Monitor render a shell and then 401 + // Workflow Runs and Guardrails Monitor render a shell and then 401 // for every non-proxy-admin role, because their page-load routes sit outside // internal_user_routes / self_managed_routes. Cost Optimization does not: // its primary call is /user/daily/activity, which every role may make, so @@ -406,7 +406,7 @@ describe("Sidebar (leftnav)", () => { mockUseAuthorized.mockReset(); }); - it("hides Workflow Runs and Memory from an internal user under Agentic", async () => { + it("shows self-service Memory and hides Workflow Runs from internal users", async () => { mockUseAuthorized.mockReturnValue(authFor("internal")); renderWithProviders(); @@ -419,7 +419,7 @@ describe("Sidebar (leftnav)", () => { expect(screen.getByText("Agents")).toBeInTheDocument(); }); expect(screen.queryByText("Workflow Runs")).not.toBeInTheDocument(); - expect(screen.queryByText("Memory")).not.toBeInTheDocument(); + expect(screen.getByText("Memory")).toBeInTheDocument(); }); // An org admin's session role is "Org Admin", which no capability list