refactor(memory): remove obsolete policy and end-to-end scaffolding

This commit is contained in:
moe-berri 2026-09-15 10:17:03 -07:00
parent eb72042a0a
commit fc8637ad86
24 changed files with 34 additions and 847 deletions

View file

@ -71,7 +71,6 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/v1/workflows/",
"/project/",
"/memory/",
"/memory/v2/",
"/mcp/",
# Control plane (see the List Endpoints + Tables standard). Every resource
# eventually moves under this prefix, so allowlist it once rather than

View file

@ -3,24 +3,23 @@ ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN IF NOT EXISTS "namespace" TEXT;
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_namespace_updated_at_idx"
ON "LiteLLM_MemoryTable"("namespace", "updated_at");
CREATE TABLE IF NOT EXISTS "LiteLLM_MemoryPolicy" (
"policy_id" TEXT NOT NULL,
"target_type" TEXT NOT NULL,
"target_id" TEXT NOT NULL,
"activation" TEXT NOT NULL,
"scope" TEXT NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_by" TEXT NOT NULL,
CONSTRAINT "LiteLLM_MemoryPolicy_pkey" PRIMARY KEY ("policy_id")
ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN IF NOT EXISTS "organization_id" TEXT;
ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN IF NOT EXISTS "owner_key_id" TEXT;
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_organization_id_user_id_updated_at_idx" ON "LiteLLM_MemoryTable"("organization_id", "user_id", "updated_at");
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_team_id_updated_at_idx" ON "LiteLLM_MemoryTable"("team_id", "updated_at");
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_owner_key_id_idx" ON "LiteLLM_MemoryTable"("owner_key_id");
CREATE TABLE IF NOT EXISTS "LiteLLM_MemoryContinuation" (
"id" TEXT NOT NULL,
"namespace" TEXT NOT NULL,
"key_id" TEXT NOT NULL,
"payload" JSONB NOT NULL,
"expires_at" TIMESTAMP(3) NOT NULL,
CONSTRAINT "LiteLLM_MemoryContinuation_pkey" PRIMARY KEY ("id")
);
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MemoryPolicy_target_type_target_id_key"
ON "LiteLLM_MemoryPolicy"("target_type", "target_id");
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryContinuation_namespace_key_id_expires_at_idx"
ON "LiteLLM_MemoryContinuation"("namespace", "key_id", "expires_at");
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,
CONSTRAINT "LiteLLM_MemoryPreference_pkey" PRIMARY KEY ("subject")
);
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryContinuation_expires_at_idx"
ON "LiteLLM_MemoryContinuation"("expires_at");

View file

@ -1,14 +0,0 @@
CREATE TABLE IF NOT EXISTS "LiteLLM_MemoryContinuation" (
"id" TEXT NOT NULL,
"namespace" TEXT NOT NULL,
"key_id" TEXT NOT NULL,
"payload" JSONB NOT NULL,
"expires_at" TIMESTAMP(3) NOT NULL,
CONSTRAINT "LiteLLM_MemoryContinuation_pkey" PRIMARY KEY ("id")
);
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryContinuation_namespace_key_id_expires_at_idx"
ON "LiteLLM_MemoryContinuation"("namespace", "key_id", "expires_at");
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryContinuation_expires_at_idx"
ON "LiteLLM_MemoryContinuation"("expires_at");

View file

@ -1 +0,0 @@
ALTER TABLE "LiteLLM_MemoryPolicy" ADD COLUMN IF NOT EXISTS "paused" BOOLEAN NOT NULL DEFAULT false;

View file

@ -1,5 +0,0 @@
ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN IF NOT EXISTS "organization_id" TEXT;
ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN IF NOT EXISTS "owner_key_id" TEXT;
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_organization_id_user_id_updated_at_idx" ON "LiteLLM_MemoryTable"("organization_id", "user_id", "updated_at");
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_team_id_updated_at_idx" ON "LiteLLM_MemoryTable"("team_id", "updated_at");
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_owner_key_id_idx" ON "LiteLLM_MemoryTable"("owner_key_id");

View file

@ -1451,26 +1451,6 @@ model LiteLLM_MemoryTable {
@@index([owner_key_id])
}
model LiteLLM_MemoryPolicy {
policy_id String @id
target_type String
target_id String
activation String
scope String
paused Boolean @default(false)
created_at DateTime @default(now())
updated_at DateTime @default(now()) @updatedAt
updated_by String
@@unique([target_type, target_id])
}
model LiteLLM_MemoryPreference {
subject String @id
enabled Boolean @default(false)
updated_at DateTime @default(now()) @updatedAt
}
model LiteLLM_MemoryContinuation {
id String @id
namespace String

View file

@ -932,9 +932,6 @@ class LiteLLMRoutes(enum.Enum):
# updating this list — the default-allow behavior covers it automatically.
admin_viewer_routes = (
[
"/memory/v2/settings",
"/memory/v2/status",
"/memory/v2/entries",
"/user/list",
"/user/available_users",
"/user/available_roles",

View file

@ -203,7 +203,7 @@ async def _assert_write_access(
)
async def is_memory_team_admin(prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool:
async def _is_team_admin_for(prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool:
"""
True if the caller is a team admin of `team_id`, or an org admin for the
team's organization. Mirrors the auth pattern used by team-management
@ -591,4 +591,3 @@ async def delete_memory(
_require_prisma = require_memory_prisma
_is_team_admin_for = is_memory_team_admin

View file

@ -70,7 +70,6 @@ class MemoryIdentity:
key_id: str | None
user_id: str | None
team_id: str | None
project_id: str | None
organization_id: str | None
read_only: bool
role: str | None = None
@ -89,7 +88,6 @@ class MemoryIdentity:
key_id=key_id,
user_id=auth.user_id,
team_id=None if dashboard else auth.team_id,
project_id=auth.project_id,
organization_id=auth.org_id,
read_only=auth.user_role
in (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY),

View file

@ -70,10 +70,10 @@ def _before(cursor: tuple[datetime, str] | None) -> Mapping[str, object]:
class MemoryStore:
def __init__(self, prisma_client: object, access: MemoryAccess, *, actor: str | None = None) -> None:
def __init__(self, prisma_client: object, access: MemoryAccess) -> None:
self.prisma_client = memory_primary_client(prisma_client)
self.access = access
self.actor = actor or access.identity.user_id or access.identity.key_id
self.actor = access.identity.user_id or access.identity.key_id
self.table = MemoryRepository(self.prisma_client).table
async def authorize(self, *, write: bool = False, require_active: bool = True) -> MemoryAccess:

View file

@ -1451,26 +1451,6 @@ model LiteLLM_MemoryTable {
@@index([owner_key_id])
}
model LiteLLM_MemoryPolicy {
policy_id String @id
target_type String
target_id String
activation String
scope String
paused Boolean @default(false)
created_at DateTime @default(now())
updated_at DateTime @default(now()) @updatedAt
updated_by String
@@unique([target_type, target_id])
}
model LiteLLM_MemoryPreference {
subject String @id
enabled Boolean @default(false)
updated_at DateTime @default(now()) @updatedAt
}
model LiteLLM_MemoryContinuation {
id String @id
namespace String

View file

@ -98,7 +98,7 @@ class LinkedSpendResetWrites:
table: BatchTable
def queue_spend_zero(self, where: Mapping[str, object]) -> None:
self.table.update_many(where=where, data={"spend": 0}) # mutable-ok: Callable type parameter syntax.
self.table.update_many(where=where, data={"spend": 0})
def queue_spend_decrement(self, where: Mapping[str, object], amount: float) -> None:
"""``decrement`` rather than a read-then-set, so spend written between the
@ -116,9 +116,7 @@ class BudgetWindowWrites:
def queue_window_advance(self, budget_id: str, budget_reset_at: datetime) -> None:
"""``update_many`` so a tier deleted between the read and the commit is a
no-op row count instead of a P2025 that aborts the whole chunk."""
self.table.update_many(
where={"budget_id": budget_id}, data={"budget_reset_at": budget_reset_at}
) # mutable-ok: Callable type parameter syntax.
self.table.update_many(where={"budget_id": budget_id}, data={"budget_reset_at": budget_reset_at})
@dataclass(frozen=True, slots=True)

View file

@ -89,12 +89,6 @@ class MemorySearch(BaseModel):
offset: int = Field(default=0, ge=0, le=10000)
class MemoryRead(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
memory_id: str = Field(min_length=1, max_length=64)
class MemoryObservation(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)

View file

@ -1451,26 +1451,6 @@ model LiteLLM_MemoryTable {
@@index([owner_key_id])
}
model LiteLLM_MemoryPolicy {
policy_id String @id
target_type String
target_id String
activation String
scope String
paused Boolean @default(false)
created_at DateTime @default(now())
updated_at DateTime @default(now()) @updatedAt
updated_by String
@@unique([target_type, target_id])
}
model LiteLLM_MemoryPreference {
subject String @id
enabled Boolean @default(false)
updated_at DateTime @default(now()) @updatedAt
}
model LiteLLM_MemoryContinuation {
id String @id
namespace String

View file

@ -90,15 +90,6 @@
- {id: mgmt.mcp_toolset.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3098", rationale: "Narrowing the tools to one entry reads back exactly that entry"}
- {id: mgmt.mcp_toolset.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:3098", fail_before_fix: proven, rationale: "An explicit null clears the stored description; the update used to drop null and keep the old value"}
- {id: mgmt.mcp_toolset.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3149", rationale: "A deleted toolset is gone by id and from the list on every replica"}
- {id: mgmt.memory_v2.gateway.capture_recall, module: mgmt, tier: P0, surface: api, assertions: [capture_recall], source: "memory/gateway.py", rationale: "Existing clients store and recall across fresh conversations and native streams"}
- {id: mgmt.memory_v2.settings.enrollment, module: mgmt, tier: P0, surface: api, assertions: [enrollment], source: "memory/policy.py", rationale: "Admin enrollment follows owners across keys; disabling stops automatic memory and preserves authorized inspection"}
- {id: mgmt.memory_v2.entries.isolation, module: mgmt, tier: P0, surface: api, assertions: [isolation], source: "memory/store.py", rationale: "Owner keys share records while unrelated users and V1 remain isolated"}
- {id: mgmt.memory_v2.settings.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "memory/management.py", rationale: "Only proxy administrators can enable memory for everyone or selected users"}
- {id: mgmt.memory_v2.entries.correction_delete, module: mgmt, tier: P0, surface: api, assertions: [correction_delete], source: "memory/store.py", rationale: "Corrections reject stale revisions and deletion removes facts from future recall"}
- {id: mgmt.memory_v2.gateway.client_tools, module: mgmt, tier: P0, surface: api, assertions: [client_tools], source: "memory/gateway.py", rationale: "Application tool calls and their continuations stay under client control"}
- {id: mgmt.memory_v2.gateway.billing, module: mgmt, tier: P0, surface: api, assertions: [billing], source: "memory/gateway.py", rationale: "Preparation and final answer each create distinct billed calls under the authenticated key, user, and team"}
- {id: mgmt.memory_v2.entries.team_permissions, module: mgmt, tier: P0, surface: api, assertions: [team_permissions], source: "memory/policy.py", rationale: "Existing team grants allow reads without writes; log grants do not confer memory access; revocation takes effect"}
- {id: mgmt.user.jwt.database_roles, module: mgmt, tier: P0, surface: api, assertions: [database_roles], source: "auth/handle_jwt.py", rationale: "User-only JWT subjects retain their seeded database roles and memberships"}
- {id: mgmt.key.jwt.viewer_denied, module: mgmt, tier: P0, surface: api, assertions: [viewer_denied], source: "auth/route_checks.py", rationale: "An admin viewer can read a key but cannot update it or change stored state"}

View file

@ -1,109 +0,0 @@
from dataclasses import dataclass
from e2e_http import NoBody, Result, unwrap
from models import (
MemoryCaptureBody,
MemoryEntriesData,
MemoryEntryData,
MemoryEntryParams,
MemorySettingsBody,
MemoryStatusData,
)
from proxy_client import ProxyClient
@dataclass(frozen=True)
class MemoryClient:
proxy: ProxyClient
def settings(self) -> MemorySettingsBody:
return unwrap(
self.proxy.transport.get(
"/memory/v2/settings",
headers=self.proxy.transport.master,
params=NoBody(),
response_type=MemorySettingsBody,
)
)
def set_settings(self, body: MemorySettingsBody, *, caller: str | None = None) -> Result[MemorySettingsBody]:
return self.proxy.transport.put(
"/memory/v2/settings",
headers=self.proxy.transport.bearer(caller) if caller else self.proxy.transport.master,
json=body,
response_type=MemorySettingsBody,
)
def read(self, key: str, memory_id: str) -> Result[MemoryEntryData]:
return self.proxy.transport.get(
f"/memory/v2/entries/{memory_id}",
headers=self.proxy.transport.bearer(key),
params=NoBody(),
response_type=MemoryEntryData,
)
def update(self, key: str, memory_id: str, body: MemoryCaptureBody) -> Result[MemoryEntryData]:
return self.proxy.transport.put(
f"/memory/v2/entries/{memory_id}",
headers=self.proxy.transport.bearer(key),
json=body,
response_type=MemoryEntryData,
)
def status(self, key: str) -> MemoryStatusData:
return unwrap(
self.proxy.transport.get(
"/memory/v2/status",
headers=self.proxy.transport.bearer(key),
params=NoBody(),
response_type=MemoryStatusData,
)
)
def entries(self, key: str, params: MemoryEntryParams = MemoryEntryParams()) -> list[MemoryEntryData]:
return unwrap(
self.proxy.transport.get(
"/memory/v2/entries",
headers=self.proxy.transport.bearer(key),
params=params,
response_type=MemoryEntriesData,
)
).root
def capture(self, key: str, body: MemoryCaptureBody) -> Result[MemoryEntryData]:
return self.proxy.transport.post(
"/memory/v2/entries",
headers=self.proxy.transport.bearer(key),
json=body,
response_type=MemoryEntryData,
)
def delete_entry(self, key: str, memory_id: str) -> Result[NoBody]:
return self.proxy.transport.delete(
f"/memory/v2/entries/{memory_id}",
headers=self.proxy.transport.bearer(key),
json=NoBody(),
response_type=NoBody,
)
def cleanup_user_entries(self, user_id: str) -> None:
while True:
page = unwrap(
self.proxy.transport.get(
"/memory/v2/entries",
headers=self.proxy.transport.master,
params=MemoryEntryParams(user_id=user_id),
response_type=MemoryEntriesData,
)
).root
if not page:
return
for entry in page:
unwrap(
self.proxy.transport.delete(
f"/memory/v2/entries/{entry.memory_id}",
headers=self.proxy.transport.master,
json=NoBody(),
response_type=NoBody,
)
)

View file

@ -1,440 +0,0 @@
import fcntl
import hashlib
import os
import tempfile
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_http import NoBody, Success, unwrap
from lifecycle import ResourceManager
from management_client import ManagementClient
from memory_client import MemoryClient
from models import (
AnthropicMessagesBody,
ChatAssistantTurn,
ChatBody,
ChatMessage,
ChatTool,
ChatToolFunction,
ChatToolResultTurn,
KeyGenerateBody,
LiteLLMParamsBody,
MemoryCaptureBody,
MemoryEntryParams,
MemoryLegacyParams,
MemoryLegacyRows,
MemoryResponsesBody,
MemorySettingsBody,
MemoryStreamEvent,
MemoryTeamPermissionBody,
MemoryWireResponse,
TeamNewBody,
UserNewBody,
)
pytestmark = pytest.mark.e2e
@dataclass(frozen=True)
class MemorySubjects:
owner: str
sibling: str
outsider: str
user_id: str
team_id: str
@dataclass(frozen=True, slots=True)
class MemoryModels:
chat: str
messages: str
@pytest.fixture
def memory_models(client: ManagementClient, resources: ResourceManager) -> MemoryModels:
def register(name: str, model: str, credential: str) -> str:
alias: Final = f"e2e-memory-{name}-{unique_marker()}"
identifier: Final = client.proxy.create_model(
alias,
LiteLLMParamsBody(
model=model,
api_key=credential,
api_base=os.environ.get("E2E_MEMORY_API_BASE"),
),
)
resources.defer(lambda: client.proxy.delete_model(identifier))
return alias
return MemoryModels(
chat=register(
"chat", os.environ.get("E2E_MEMORY_CHAT_MODEL", "openai/gpt-5.6-sol"), "os.environ/OPENAI_API_KEY"
),
messages=register(
"messages",
os.environ.get("E2E_MEMORY_MESSAGES_MODEL", "anthropic/claude-haiku-4-5"),
"os.environ/ANTHROPIC_API_KEY",
),
)
@pytest.fixture
def memory(client: ManagementClient) -> Iterator[MemoryClient]:
# Serialize global configuration changes across local pytest workers.
with (Path(tempfile.gettempdir()) / "litellm-memory-v2-e2e.lock").open("w") as lock:
fcntl.flock(lock, fcntl.LOCK_EX)
memory = MemoryClient(client.proxy)
original = memory.settings()
try:
yield memory
finally:
unwrap(memory.set_settings(original))
fcntl.flock(lock, fcntl.LOCK_UN)
@pytest.fixture
def subjects(client: ManagementClient, memory: MemoryClient, resources: ResourceManager) -> MemorySubjects:
marker: Final = unique_marker()
owner: Final = client.create_user(
UserNewBody(
user_email=f"memory-owner-{marker}@example.invalid", user_role="internal_user", auto_create_key=False
)
)
resources.defer(lambda: client.delete_user_strict(owner))
other: Final = client.create_user(
UserNewBody(
user_email=f"memory-other-{marker}@example.invalid", user_role="internal_user", auto_create_key=False
)
)
resources.defer(lambda: client.delete_user_strict(other))
team: Final = client.create_team(TeamNewBody(team_alias=f"memory-{marker}"))
resources.defer(lambda: client.delete_team(team))
client.add_team_member(team, owner)
client.add_team_member(team, other)
def create_key(user: str) -> str:
key: Final = unwrap(
client.generate_key(KeyGenerateBody(user_id=user, team_id=team, models=[], max_parallel_requests=1))
).key
resources.defer(lambda: client.delete_key_strict(key))
return key
keys: Final = tuple(create_key(user) for user in (owner, owner, other))
unwrap(memory.set_settings(MemorySettingsBody(enabled=True, everyone=False, user_ids=[owner, other])))
resources.defer(lambda: memory.cleanup_user_entries(owner))
resources.defer(lambda: memory.cleanup_user_entries(other))
return MemorySubjects(owner=keys[0], sibling=keys[1], outsider=keys[2], user_id=owner, team_id=team)
def _fact(marker: str) -> MemoryCaptureBody:
return MemoryCaptureBody(
key=f"release-{marker}",
title="Demo project codename",
content=f"The demo project codename is {marker}",
evidence="Synthetic fact supplied by the authenticated test user",
)
def _assert_denied(result: object) -> None:
assert not isinstance(result, Success), "An unauthorized memory operation succeeded"
assert "403" in str(result) or "unauthorized" in str(result).lower() or "404" in str(result), result
class TestMemoryV2:
@pytest.mark.covers("mgmt.memory_v2.gateway.capture_recall")
@pytest.mark.parametrize("endpoint", ["chat", "responses", "messages"])
@pytest.mark.parametrize("stream", [False, True])
def test_gateway_stores_and_recalls_without_client_memory_tools(
self,
client: ManagementClient,
memory: MemoryClient,
subjects: MemorySubjects,
memory_models: MemoryModels,
endpoint: str,
stream: bool,
) -> None:
marker: Final = f"copper-{unique_marker()}"
seed: Final = unwrap(
client.proxy.chat(
subjects.owner,
ChatBody(
model=memory_models.chat,
max_tokens=1200,
messages=[
ChatMessage(
role="user",
content=f"Remember this durable preference for future conversations: my demo project codename is {marker}. Confirm briefly.",
)
],
),
)
)
assert seed.choices
stored: Final = memory.entries(subjects.owner)
assert any(marker in entry.content for entry in stored), stored
prompt: Final = "What is my demo project codename? Return the exact word only."
model: Final = memory_models.messages if endpoint == "messages" else memory_models.chat
body: Final = (
AnthropicMessagesBody(
model=model, messages=[ChatMessage(role="user", content=prompt)], max_tokens=1200, stream=stream
)
if endpoint == "messages"
else MemoryResponsesBody(model=model, input=prompt, stream=stream)
if endpoint == "responses"
else ChatBody(
model=model, messages=[ChatMessage(role="user", content=prompt)], max_tokens=1200, stream=stream
)
)
path: Final = {"messages": "/v1/messages", "responses": "/v1/responses", "chat": "/v1/chat/completions"}[
endpoint
]
response: Final = client.proxy.transport.send(
path, headers=client.proxy.transport.bearer(subjects.owner), json=body, stream=stream
)
assert response.status_code == 200, response.body
assert response.stream_error is None, response.stream_error
output: Final = (
"".join(MemoryStreamEvent.model_validate_json(event).text for event in response.stream_events)
if stream
else response.body
)
assert marker in output, output
if stream:
assert response.is_streaming
assert response.chunks > 1
events: Final = tuple(MemoryStreamEvent.model_validate_json(event) for event in response.stream_events)
assert not any(event.has_memory_tools for event in events)
if endpoint == "responses":
assert all(event.response.instructions is None for event in events if event.response)
else:
public: Final = MemoryWireResponse.model_validate_json(response.body)
assert not public.has_memory_tools
if endpoint == "responses":
assert public.instructions is None
assert any(marker in entry.content for entry in memory.entries(subjects.sibling))
assert memory.entries(subjects.outsider) == []
@pytest.mark.covers("mgmt.memory_v2.settings.enrollment")
def test_admin_enrollment_follows_user_and_disable_preserves_dashboard(
self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects, memory_models: MemoryModels
) -> None:
unwrap(memory.set_settings(MemorySettingsBody()))
assert not memory.status(subjects.owner).active
marker = f"unsaved-{unique_marker()}"
unwrap(
client.proxy.chat(
subjects.owner,
ChatBody(
model=memory_models.chat,
max_tokens=100,
messages=[
ChatMessage(
role="user", content=f"Remember that my verification word is {marker}. Confirm briefly."
)
],
),
)
)
assert memory.entries(subjects.owner) == []
unwrap(memory.set_settings(MemorySettingsBody(enabled=True, everyone=False, user_ids=[subjects.user_id])))
assert memory.status(subjects.owner).active and memory.status(subjects.sibling).active
assert not memory.status(subjects.outsider).active
saved = unwrap(memory.capture(subjects.owner, _fact(unique_marker())))
unwrap(memory.set_settings(MemorySettingsBody()))
assert not memory.status(subjects.owner).active
assert unwrap(memory.read(subjects.sibling, saved.memory_id)).content == saved.content
_assert_denied(memory.capture(subjects.owner, _fact(unique_marker())))
@pytest.mark.covers("mgmt.memory_v2.entries.isolation")
def test_owner_keys_share_but_other_users_and_legacy_api_do_not(
self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects
) -> None:
saved = unwrap(memory.capture(subjects.owner, _fact(unique_marker())))
assert memory.entries(subjects.sibling)[0].memory_id == saved.memory_id
assert memory.entries(subjects.outsider) == []
assert memory.entries(subjects.outsider, MemoryEntryParams(user_id=subjects.user_id)) == []
_assert_denied(memory.read(subjects.outsider, saved.memory_id))
_assert_denied(memory.delete_entry(subjects.outsider, saved.memory_id))
legacy = client.proxy.transport.get(
"/v1/memory",
headers=client.proxy.transport.bearer(subjects.sibling),
params=MemoryLegacyParams(),
response_type=MemoryLegacyRows,
)
if isinstance(legacy, Success):
assert saved.memory_id not in [row.memory_id for row in legacy.data.memories]
assert memory.entries(subjects.owner)[0].content == saved.content
@pytest.mark.covers("mgmt.memory_v2.settings.admin_only")
def test_members_cannot_enable_memory(self, memory: MemoryClient, subjects: MemorySubjects) -> None:
before = memory.settings()
_assert_denied(memory.set_settings(MemorySettingsBody(enabled=True), caller=subjects.owner))
assert memory.settings() == before
@pytest.mark.covers("mgmt.memory_v2.entries.team_permissions")
def test_delegated_team_reads_allow_recall_but_not_edit_and_can_be_revoked(
self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects, memory_models: MemoryModels
) -> None:
marker = f"team-{unique_marker()}"
saved = unwrap(memory.capture(subjects.owner, _fact(marker)))
assert memory.entries(subjects.outsider) == []
for permissions, visible in ((["/spend/logs"], False), (["/memory/v2/entries"], True), ([], False)):
unwrap(
client.proxy.transport.post(
"/team/permissions_update",
headers=client.proxy.transport.master,
json=MemoryTeamPermissionBody(team_id=subjects.team_id, team_member_permissions=permissions),
response_type=NoBody,
)
)
entries = memory.entries(subjects.outsider, MemoryEntryParams(team_id=subjects.team_id))
assert bool(entries) is visible
if visible:
assert entries[0].memory_id == saved.memory_id and not entries[0].can_edit
assert unwrap(memory.read(subjects.outsider, saved.memory_id)).content == saved.content
_assert_denied(memory.update(subjects.outsider, saved.memory_id, _fact(unique_marker())))
_assert_denied(memory.delete_entry(subjects.outsider, saved.memory_id))
recalled = unwrap(
client.proxy.chat(
subjects.outsider,
ChatBody(
model=memory_models.chat,
max_tokens=1200,
messages=[
ChatMessage(
role="user",
content="Search the team's memories for the demo project codename and return it exactly.",
)
],
),
)
)
assert marker in recalled.model_dump_json()
else:
_assert_denied(memory.read(subjects.outsider, saved.memory_id))
@pytest.mark.covers("mgmt.memory_v2.entries.correction_delete")
def test_corrections_require_current_revision_and_deleted_memory_is_not_recalled(
self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects, memory_models: MemoryModels
) -> None:
original: Final = _fact(unique_marker())
saved: Final = unwrap(memory.capture(subjects.owner, original))
corrected_word: Final = f"corrected-{unique_marker()}"
correction: Final = MemoryCaptureBody(
key=original.key,
title=original.title,
content=f"The demo project codename is {corrected_word}",
evidence="The user corrected the earlier word",
expected_revision=saved.updated_at,
)
updated: Final = unwrap(memory.capture(subjects.owner, correction))
assert updated.memory_id == saved.memory_id
stale: Final = memory.capture(
subjects.owner, original.model_copy(update={"expected_revision": saved.updated_at})
)
assert not isinstance(stale, Success)
assert "409" in str(stale), stale
assert memory.entries(subjects.owner)[0].content == correction.content
unwrap(memory.delete_entry(subjects.owner, saved.memory_id))
assert memory.entries(subjects.owner) == []
response: Final = unwrap(
client.proxy.chat(
subjects.owner,
ChatBody(
model=memory_models.chat,
max_tokens=200,
messages=[
ChatMessage(
role="user", content="What is my demo project codename? If it is unknown, say unknown."
)
],
),
)
)
assert corrected_word not in response.model_dump_json()
@pytest.mark.covers("mgmt.memory_v2.gateway.client_tools")
def test_client_tool_call_and_continuation_remain_owned_by_client(
self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects, memory_models: MemoryModels
) -> None:
marker: Final = unique_marker()
unwrap(memory.capture(subjects.owner, _fact(marker)))
tool: Final = ChatTool(
function=ChatToolFunction(
name="verify_release",
description="Verify a release using the user's verification word",
parameters={"type": "object", "properties": {"word": {"type": "string"}}, "required": ["word"]},
)
)
prompt: Final = ChatMessage(
role="user",
content="Use verify_release with my demo project codename. After the tool returns, save its verification result in memory with the tool as your evidence, then report it.",
)
response: Final = unwrap(
client.proxy.chat(
subjects.owner,
ChatBody(
model=memory_models.chat, messages=[prompt], tools=[tool], tool_choice="required", max_tokens=1200
),
)
)
message: Final = response.choices[0].message
assert message is not None
calls: Final = message.tool_calls
assert calls and len(calls) == 1
call: Final = calls[0]
assert call.function.name == "verify_release"
assert marker in (call.function.arguments or "")
assert call.id
result_marker: Final = f"verified-{unique_marker()}"
followup: Final = unwrap(
client.proxy.chat(
subjects.owner,
ChatBody(
model=memory_models.chat,
max_tokens=1200,
tools=[tool],
messages=[
prompt,
ChatAssistantTurn(
content=message.content, reasoning_content=message.reasoning_content, tool_calls=calls
),
ChatToolResultTurn(tool_call_id=call.id, content=result_marker),
],
),
)
)
assert result_marker in followup.model_dump_json()
assert any(result_marker in entry.content for entry in memory.entries(subjects.owner))
@pytest.mark.covers("mgmt.memory_v2.gateway.billing")
def test_memory_tool_rounds_are_charged_once_to_the_calling_key(
self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects, memory_models: MemoryModels
) -> None:
marker: Final = unique_marker()
response: Final = unwrap(
client.proxy.chat(
subjects.owner,
ChatBody(
model=memory_models.chat,
max_tokens=1200,
messages=[
ChatMessage(
role="user", content=f"Remember that my demo project codename is {marker}. Confirm briefly."
)
],
),
)
)
assert response.choices
assert any(marker in row.content for row in memory.entries(subjects.owner))
rows: Final = client.proxy.poll_logs_for_key(subjects.owner, min_rows=2)
assert 2 <= len(rows) <= 8, rows
assert len({row.request_id for row in rows}) == len(rows), rows
assert all(row.api_key == hashlib.sha256(subjects.owner.encode()).hexdigest() for row in rows), rows
assert all(row.user == subjects.user_id and row.team_id == subjects.team_id for row in rows), rows
assert all(row.prompt_tokens and row.completion_tokens for row in rows), rows
assert all(row.spend is not None and row.spend > 0 for row in rows), rows

View file

@ -69,7 +69,6 @@ class ObjectPermission(BaseModel):
class KeyGenerateBody(BaseModel):
max_parallel_requests: int | None = None
models: list[str] = []
duration: str | None = None
max_budget: float | None = None
@ -1327,154 +1326,3 @@ class ReadinessDetailsResponse(ReadinessResponse):
litellm_version: str | None = None
success_callbacks: list[str] = []
class MemorySettingsBody(BaseModel):
enabled: bool = False
everyone: bool = True
user_ids: list[str] = []
class MemoryStatusData(BaseModel):
active: bool
enabled: bool
user_id: str | None = None
team_ids: list[str] = []
admin_view: bool = False
class MemoryEntryParams(BaseModel):
query: str = ""
limit: int = 20
user_id: str | None = None
team_id: str | None = None
offset: int = 0
before_updated_at: str | None = None
before_memory_id: str | None = None
class MemoryTeamPermissionBody(BaseModel):
team_id: str
team_member_permissions: list[str]
class MemoryCaptureBody(BaseModel):
key: str
title: str
content: str
evidence: str
expected_revision: str | None = None
class MemoryEntryData(BaseModel):
user_id: str | None = None
team_id: str | None = None
can_edit: bool = False
actor: str | None = None
memory_id: str
key: str
title: str
content: str
evidence: str
updated_at: str
class MemoryEntriesData(RootModel[list[MemoryEntryData]]):
pass
class MemoryLegacyParams(BaseModel):
key_prefix: str = "memory-v2:"
page: int = 1
page_size: int = 500
class MemoryLegacyRow(BaseModel):
memory_id: str
key: str
user_id: str | None = None
class MemoryLegacyRows(BaseModel):
memories: list[MemoryLegacyRow]
total: int
class MemoryResponsesBody(BaseModel):
model: str
input: str
stream: bool
max_output_tokens: int = 1200
store: bool = False
cache: dict[str, bool] = {"no-cache": True}
class MemoryWireTool(BaseModel):
name: str | None = None
function: ToolCallFunction = ToolCallFunction()
@property
def is_gateway_memory(self) -> bool:
return (self.name or self.function.name or "").startswith("litellm_memory_")
class MemoryStreamDelta(BaseModel):
content: str | None = None
text: str | None = None
tool_calls: tuple[MemoryWireTool, ...] | None = None
class MemoryStreamChoice(BaseModel):
delta: MemoryStreamDelta = MemoryStreamDelta()
message: MemoryStreamDelta = MemoryStreamDelta()
class MemoryWireResponse(BaseModel):
instructions: str | None = None
tools: tuple[MemoryWireTool, ...] = ()
output: tuple[MemoryWireTool, ...] = ()
content: tuple[MemoryWireTool, ...] = ()
choices: tuple[MemoryStreamChoice, ...] = ()
@property
def has_memory_tools(self) -> bool:
return any(
tool.is_gateway_memory
for tool in (
*self.tools,
*self.output,
*self.content,
*(
tool
for choice in self.choices
for part in (choice.delta, choice.message)
for tool in part.tool_calls or ()
),
)
)
class MemoryStreamEvent(BaseModel):
delta: MemoryStreamDelta | str | None = None
choices: list[MemoryStreamChoice] = []
response: MemoryWireResponse | None = None
item: MemoryWireTool = MemoryWireTool()
content_block: MemoryWireTool = MemoryWireTool()
@property
def has_memory_tools(self) -> bool:
return bool(
self.response
and self.response.has_memory_tools
or self.item.is_gateway_memory
or self.content_block.is_gateway_memory
or any(tool.is_gateway_memory for choice in self.choices for tool in choice.delta.tool_calls or ())
)
@property
def text(self) -> str:
if isinstance(self.delta, str):
return self.delta
if self.delta:
return self.delta.text or self.delta.content or ""
return "".join(choice.delta.content or "" for choice in self.choices)

View file

@ -305,7 +305,6 @@ CONTROL_PLANE_PREFIXES: tuple[str, ...] = (
"/global",
"/config",
"/guardrails",
"/memory/v2/settings",
"/openapi.json",
)

View file

@ -1182,10 +1182,8 @@ async def test_end_user_jwt_auth(monkeypatch):
"llm_router",
router,
)
# 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)
memory_db.db.litellm_config.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

View file

@ -503,10 +503,8 @@ 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")
# 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)
memory_db.db.litellm_config.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)

View file

@ -27,7 +27,7 @@ from litellm.proxy.memory.store import MemoryStore
from litellm.types.memory_v2 import MemoryCapture, MemorySearch, MemorySettings
_NOW: Final = datetime(2026, 9, 12, tzinfo=timezone.utc)
_IDENTITY: Final = MemoryIdentity("a" * 64, "owner", "team", "project", "org", False)
_IDENTITY: Final = MemoryIdentity("a" * 64, "owner", "team", "org", False)
_SETTINGS: Final = MemorySettings(enabled=True)
_CAPTURE: Final = MemoryCapture(key="demo", title="Demo", content="Use port 8347", evidence="User selected this port")
@ -162,7 +162,7 @@ async def test_response_deletion_preserves_auth_paths_and_retry_state(prisma_edg
"query_string": b"api-version=test",
}
)
identity = MemoryIdentity("a" * 64, "owner", "team", "project", "org", outcome == "readonly")
identity = MemoryIdentity("a" * 64, "owner", "team", "org", outcome == "readonly")
if outcome in ("upstream_error", "readonly"):
with pytest.raises(HTTPException) as exc:
await serve_memory_response(
@ -194,7 +194,7 @@ async def test_flat_enrollment_follows_the_user_across_keys(prisma_edge: MagicMo
param_value=MemorySettings(enabled=True, everyone=False, user_ids=("owner",)).model_dump()
)
assert (await resolve_memory_access(prisma_edge, _IDENTITY)).active
sibling = MemoryIdentity("b" * 64, "owner", "team", "project", "org", False)
sibling = MemoryIdentity("b" * 64, "owner", "team", "org", False)
assert (await resolve_memory_access(prisma_edge, sibling)).active
config.return_value = SimpleNamespace(
param_value=MemorySettings(enabled=True, everyone=False, user_ids=("other",)).model_dump()
@ -209,7 +209,7 @@ async def test_store_rechecks_access_before_writing(prisma_edge: MagicMock, chan
if change == "missing":
prisma_edge.db.litellm_config.find_unique.return_value = None
elif change == "readonly":
identity = MemoryIdentity("a" * 64, "owner", "team", "project", "org", True)
identity = MemoryIdentity("a" * 64, "owner", "team", "org", True)
else:
config = MemorySettings(enabled=change != "disabled", everyone=False, user_ids=("other",))
prisma_edge.db.litellm_config.find_unique.return_value = SimpleNamespace(param_value=config.model_dump())
@ -377,7 +377,7 @@ def request() -> Request:
@pytest.mark.asyncio
async def test_read_only_injection_and_forced_no_tools_do_not_request_reflection(prisma_edge: MagicMock) -> None:
read_only: Final = MemoryIdentity("a" * 64, "owner", "team", "project", "org", True)
read_only: Final = MemoryIdentity("a" * 64, "owner", "team", "org", True)
loop: Final = GatewayMemoryLoop(
FastAPI(),
request(),
@ -468,7 +468,7 @@ async def test_restore_preserves_hidden_memory_tool_results_and_client_cache_mar
*object_items(replacement[1]["content"]),
*object_items(items[2]["content"]),
)
sibling: Final = MemoryIdentity("b" * 64, "owner", "team", "project", "org", False)
sibling: Final = MemoryIdentity("b" * 64, "owner", "team", "org", False)
assert MemoryContinuations(store(prisma_edge, sibling), "anthropic_messages").identifier(
anchor
) != continuations.identifier(anchor)
@ -1085,7 +1085,7 @@ async def test_continuation_quota_shares_namespace_lock_across_keys_and_allows_r
prisma_edge: MagicMock,
) -> None:
prisma_edge.db.query_raw.return_value = [{"key_count": 255, "bytes": 32 * 1024 * 1024}]
other_key = MemoryIdentity("b" * 64, "owner", "team", "project", "org", False)
other_key = MemoryIdentity("b" * 64, "owner", "team", "org", False)
for identity in (_IDENTITY, other_key):
continuations = MemoryContinuations(MemoryStore(prisma_edge, access_for(identity)), "aresponses")
await continuations.save_many((("response", MemoryContinuation(replaces=1, response={"text": "é漢字"})),))

View file

@ -4298,7 +4298,7 @@ class TestDisconnectGatherCleanup:
await processing_obj.base_process_llm_request(
request=self._disconnect_request(),
fastapi_response=MagicMock(),
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
user_api_key_dict=ProxyUserAPIKeyAuth(),
proxy_logging_obj=mock_proxy_logging,
general_settings={"cancel_on_disconnect": True},
proxy_config=MagicMock(spec=ProxyConfig),

View file

@ -33,8 +33,6 @@ export interface DeletedKeysResponse {
}
export interface KeyListCallOptions {
includeTeamKeys?: boolean;
includeCreatedByKeys?: boolean;
organizationID?: string | null;
teamID?: string | null;
projectID?: string | null;
@ -73,8 +71,8 @@ const keyListCall = async (accessToken: string, page: number, pageSize: number,
expand: options.expand,
status: options.status,
return_full_object: "true",
include_team_keys: options.includeTeamKeys ?? true,
include_created_by_keys: options.includeCreatedByKeys ?? true,
include_team_keys: "true",
include_created_by_keys: "true",
// Opt into substring matching so the admin key-list search box keeps
// matching partial user_id/key_alias. /key/list is exact by default.
substring_matching: "true",