From 4010f471c8ab5f518e2e20a78f0ea0d66faa7df5 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 14 Sep 2026 13:52:18 -0700 Subject: [PATCH] fix(memory): preserve Claude directives and bound continuation storage --- deploy/memory-pilot/pilot.py | 5 +- .../prompt_templates/server_tools.py | 19 +- litellm/proxy/_lazy_features.py | 5 + litellm/proxy/_lazy_openapi_snapshot.json | 1074 +++++++++++++++++ litellm/proxy/memory/continuation.py | 33 +- litellm/proxy/memory/gateway.py | 7 +- litellm/proxy/proxy_server.py | 2 - .../proxy/memory/test_memory_v2_boundaries.py | 102 ++ .../proxy/memory/test_memory_v2_protocols.py | 17 + 9 files changed, 1242 insertions(+), 22 deletions(-) diff --git a/deploy/memory-pilot/pilot.py b/deploy/memory-pilot/pilot.py index fc006a98ca4..c328115e3f7 100644 --- a/deploy/memory-pilot/pilot.py +++ b/deploy/memory-pilot/pilot.py @@ -1,7 +1,6 @@ """An isolated office pilot that preserves upstream gateway credentials.""" import asyncio -import hashlib import os import secrets from collections.abc import Callable @@ -19,7 +18,7 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.prompt_templates.server_tool_responses import object_value from litellm.llms.custom_httpx.http_handler import get_async_httpx_client -from litellm.proxy._types import UI_TEAM_ID, UserAPIKeyAuth +from litellm.proxy._types import UI_TEAM_ID, UserAPIKeyAuth, hash_token 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.proxy.memory.transport import in_gateway_round @@ -101,7 +100,7 @@ class PilotGateway: if prisma_client is None: await JSONResponse({"error": "Pilot database unavailable"}, status_code=503)(scope, receive, send) return - digest: Final = hashlib.sha256(credential.encode()).hexdigest() + digest: Final = hash_token(credential) tokens: Final = VerificationTokenRepository(prisma_client) local_key: Final = await tokens.find_by_id(digest) if local_key and local_key.team_id == UI_TEAM_ID: diff --git a/litellm/litellm_core_utils/prompt_templates/server_tools.py b/litellm/litellm_core_utils/prompt_templates/server_tools.py index 47de79031a1..ce8e8114eb2 100644 --- a/litellm/litellm_core_utils/prompt_templates/server_tools.py +++ b/litellm/litellm_core_utils/prompt_templates/server_tools.py @@ -137,16 +137,33 @@ def _tool_name(tool: object) -> object: return _OBJECT.validate_python(function).get("name") if isinstance(function, dict) else definition.get("name") +def trailing_system_messages(data: Mapping[str, object], route: ServerToolRoute) -> int: + if route != "anthropic_messages": + return 0 + messages: Final = _items(data.get("messages")) + return next( + ( + index + for index, message in enumerate(reversed(messages)) + if not isinstance(message, dict) or _OBJECT.validate_python(message).get("role") != "system" + ), + len(messages), + ) + + def append_server_reference(data: Mapping[str, object], route: ServerToolRoute, reference: str) -> Mapping[str, object]: field: Final = "input" if route == "aresponses" else "messages" + messages: Final = _items(data.get(field)) + insertion: Final = len(messages) - trailing_system_messages(data, route) return { # mutable-ok: Provider wire format requires native JSON containers. **data, field: [ # mutable-ok: Provider wire format requires native JSON containers. - *_items(data.get(field)), + *messages[:insertion], { # mutable-ok: Provider wire format requires native JSON containers. "role": "user", "content": reference, }, + *messages[insertion:], ], } diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index dd1180b30ad..c7a8b4e0967 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -56,6 +56,11 @@ class LazyFeature: LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( + LazyFeature( + name="memory_v2", + module_path="litellm.proxy.memory.management", + path_prefixes=("/v2/memory",), + ), LazyFeature( name="guardrails", module_path="litellm.proxy.guardrails.guardrail_endpoints", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 7a110eff080..7c98f4f0ea5 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -31265,6 +31265,1080 @@ } } }, + "memory_v2": { + "components": { + "schemas": { + "HTTPValidationError": { + "properties": { + "detail": { + "items": { + "$ref": "#/components/schemas/ValidationError" + }, + "title": "Detail", + "type": "array" + } + }, + "title": "HTTPValidationError", + "type": "object" + }, + "MemoryCapture": { + "additionalProperties": false, + "properties": { + "certainty": { + "default": "observed", + "enum": [ + "user_stated", + "observed", + "inferred" + ], + "title": "Certainty", + "type": "string" + }, + "content": { + "maxLength": 8000, + "minLength": 1, + "title": "Content", + "type": "string" + }, + "evidence": { + "maxLength": 2000, + "minLength": 1, + "title": "Evidence", + "type": "string" + }, + "expected_revision": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Expected Revision" + }, + "key": { + "maxLength": 160, + "minLength": 1, + "pattern": "^[a-zA-Z0-9_.-]+$", + "title": "Key", + "type": "string" + }, + "kind": { + "default": "context", + "enum": [ + "workflow", + "decision", + "correction", + "learning", + "context", + "disagreement" + ], + "title": "Kind", + "type": "string" + }, + "scope": { + "default": "", + "maxLength": 200, + "title": "Scope", + "type": "string" + }, + "source": { + "default": "", + "maxLength": 1000, + "title": "Source", + "type": "string" + }, + "title": { + "maxLength": 200, + "minLength": 1, + "title": "Title", + "type": "string" + }, + "when_to_use": { + "default": "", + "maxLength": 700, + "title": "When To Use", + "type": "string" + } + }, + "required": [ + "key", + "title", + "content", + "evidence" + ], + "title": "MemoryCapture", + "type": "object" + }, + "MemoryEntry": { + "properties": { + "actor": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Actor" + }, + "actor_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Actor Name" + }, + "certainty": { + "default": "observed", + "enum": [ + "user_stated", + "observed", + "inferred" + ], + "title": "Certainty", + "type": "string" + }, + "content": { + "title": "Content", + "type": "string" + }, + "created_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Created At" + }, + "evidence": { + "title": "Evidence", + "type": "string" + }, + "key": { + "title": "Key", + "type": "string" + }, + "kind": { + "default": "context", + "enum": [ + "workflow", + "decision", + "correction", + "learning", + "context", + "disagreement" + ], + "title": "Kind", + "type": "string" + }, + "memory_id": { + "title": "Memory Id", + "type": "string" + }, + "scope": { + "default": "", + "title": "Scope", + "type": "string" + }, + "source": { + "default": "", + "title": "Source", + "type": "string" + }, + "title": { + "title": "Title", + "type": "string" + }, + "updated_at": { + "format": "date-time", + "title": "Updated At", + "type": "string" + }, + "when_to_use": { + "default": "", + "title": "When To Use", + "type": "string" + } + }, + "required": [ + "memory_id", + "key", + "title", + "content", + "evidence", + "updated_at" + ], + "title": "MemoryEntry", + "type": "object" + }, + "MemoryPolicy": { + "additionalProperties": false, + "properties": { + "activation": { + "enum": [ + "disabled", + "opt_in", + "automatic" + ], + "title": "Activation", + "type": "string" + }, + "policy_id": { + "title": "Policy Id", + "type": "string" + }, + "scope": { + "default": "user", + "enum": [ + "key", + "user", + "team", + "project", + "organization" + ], + "title": "Scope", + "type": "string" + }, + "target_id": { + "maxLength": 256, + "minLength": 1, + "title": "Target Id", + "type": "string" + }, + "target_type": { + "enum": [ + "gateway", + "organization", + "team", + "project", + "user", + "key" + ], + "title": "Target Type", + "type": "string" + }, + "updated_at": { + "format": "date-time", + "title": "Updated At", + "type": "string" + }, + "updated_by": { + "title": "Updated By", + "type": "string" + } + }, + "required": [ + "target_type", + "target_id", + "activation", + "policy_id", + "updated_at", + "updated_by" + ], + "title": "MemoryPolicy", + "type": "object" + }, + "MemoryPolicyInput": { + "additionalProperties": false, + "properties": { + "activation": { + "enum": [ + "disabled", + "opt_in", + "automatic" + ], + "title": "Activation", + "type": "string" + }, + "scope": { + "default": "user", + "enum": [ + "key", + "user", + "team", + "project", + "organization" + ], + "title": "Scope", + "type": "string" + }, + "target_id": { + "maxLength": 256, + "minLength": 1, + "title": "Target Id", + "type": "string" + }, + "target_type": { + "enum": [ + "gateway", + "organization", + "team", + "project", + "user", + "key" + ], + "title": "Target Type", + "type": "string" + } + }, + "required": [ + "target_type", + "target_id", + "activation" + ], + "title": "MemoryPolicyInput", + "type": "object" + }, + "MemoryPreference": { + "additionalProperties": false, + "properties": { + "enabled": { + "title": "Enabled", + "type": "boolean" + } + }, + "required": [ + "enabled" + ], + "title": "MemoryPreference", + "type": "object" + }, + "MemoryStatus": { + "properties": { + "activation": { + "enum": [ + "disabled", + "opt_in", + "automatic" + ], + "title": "Activation", + "type": "string" + }, + "active": { + "title": "Active", + "type": "boolean" + }, + "opted_in": { + "title": "Opted In", + "type": "boolean" + }, + "policy_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Policy Id" + }, + "scope": { + "anyOf": [ + { + "enum": [ + "key", + "user", + "team", + "project", + "organization" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Scope" + }, + "user_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Id" + }, + "user_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Name" + } + }, + "required": [ + "active", + "activation", + "scope", + "opted_in", + "policy_id" + ], + "title": "MemoryStatus", + "type": "object" + }, + "ValidationError": { + "properties": { + "ctx": { + "title": "Context", + "type": "object" + }, + "input": { + "title": "Input" + }, + "loc": { + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "integer" + } + ] + }, + "title": "Location", + "type": "array" + }, + "msg": { + "title": "Message", + "type": "string" + }, + "type": { + "title": "Error Type", + "type": "string" + } + }, + "required": [ + "loc", + "msg", + "type" + ], + "title": "ValidationError", + "type": "object" + } + } + }, + "paths": { + "/v2/memory/entries": { + "get": { + "operationId": "list_entries_v2_memory_entries_get", + "parameters": [ + { + "in": "query", + "name": "query", + "required": false, + "schema": { + "default": "", + "maxLength": 500, + "title": "Query", + "type": "string" + } + }, + { + "in": "query", + "name": "limit", + "required": false, + "schema": { + "default": 20, + "maximum": 20, + "minimum": 1, + "title": "Limit", + "type": "integer" + } + }, + { + "in": "query", + "name": "offset", + "required": false, + "schema": { + "default": 0, + "minimum": 0, + "title": "Offset", + "type": "integer" + } + }, + { + "in": "query", + "name": "key_id", + "required": false, + "schema": { + "anyOf": [ + { + "pattern": "^[a-f0-9]{64}$", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key Id" + } + }, + { + "in": "query", + "name": "before_updated_at", + "required": false, + "schema": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Before Updated At" + } + }, + { + "in": "query", + "name": "before_memory_id", + "required": false, + "schema": { + "anyOf": [ + { + "maxLength": 128, + "minLength": 1, + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Before Memory Id" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "items": { + "$ref": "#/components/schemas/MemoryEntry" + }, + "title": "Response List Entries V2 Memory Entries Get", + "type": "array" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "List Entries", + "tags": [ + "memory_v2" + ] + }, + "post": { + "operationId": "capture_entry_v2_memory_entries_post", + "parameters": [ + { + "in": "query", + "name": "key_id", + "required": false, + "schema": { + "anyOf": [ + { + "pattern": "^[a-f0-9]{64}$", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key Id" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MemoryCapture" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MemoryEntry" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Capture Entry", + "tags": [ + "memory_v2" + ] + } + }, + "/v2/memory/entries/{memory_id}": { + "delete": { + "operationId": "delete_entry_v2_memory_entries__memory_id__delete", + "parameters": [ + { + "in": "path", + "name": "memory_id", + "required": true, + "schema": { + "title": "Memory Id", + "type": "string" + } + }, + { + "in": "query", + "name": "key_id", + "required": false, + "schema": { + "anyOf": [ + { + "pattern": "^[a-f0-9]{64}$", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key Id" + } + } + ], + "responses": { + "204": { + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Delete Entry", + "tags": [ + "memory_v2" + ] + } + }, + "/v2/memory/policies": { + "get": { + "operationId": "list_policies_v2_memory_policies_get", + "parameters": [ + { + "in": "query", + "name": "target_type", + "required": false, + "schema": { + "anyOf": [ + { + "enum": [ + "gateway", + "organization", + "team", + "project", + "user", + "key" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Target Type" + } + }, + { + "in": "query", + "name": "target_id", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Target Id" + } + }, + { + "in": "query", + "name": "offset", + "required": false, + "schema": { + "default": 0, + "minimum": 0, + "title": "Offset", + "type": "integer" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "items": { + "$ref": "#/components/schemas/MemoryPolicy" + }, + "title": "Response List Policies V2 Memory Policies Get", + "type": "array" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "List Policies", + "tags": [ + "memory_v2" + ] + }, + "put": { + "operationId": "set_policy_v2_memory_policies_put", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MemoryPolicyInput" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MemoryPolicy" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Set Policy", + "tags": [ + "memory_v2" + ] + } + }, + "/v2/memory/policies/{policy_id}": { + "delete": { + "operationId": "delete_policy_v2_memory_policies__policy_id__delete", + "parameters": [ + { + "in": "path", + "name": "policy_id", + "required": true, + "schema": { + "title": "Policy Id", + "type": "string" + } + } + ], + "responses": { + "204": { + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Delete Policy", + "tags": [ + "memory_v2" + ] + } + }, + "/v2/memory/preference": { + "get": { + "operationId": "get_preference_v2_memory_preference_get", + "parameters": [ + { + "in": "query", + "name": "key_id", + "required": false, + "schema": { + "anyOf": [ + { + "pattern": "^[a-f0-9]{64}$", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key Id" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MemoryPreference" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Preference", + "tags": [ + "memory_v2" + ] + }, + "put": { + "operationId": "set_preference_v2_memory_preference_put", + "parameters": [ + { + "in": "query", + "name": "key_id", + "required": false, + "schema": { + "anyOf": [ + { + "pattern": "^[a-f0-9]{64}$", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key Id" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MemoryPreference" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MemoryPreference" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Set Preference", + "tags": [ + "memory_v2" + ] + } + }, + "/v2/memory/status": { + "get": { + "operationId": "get_status_v2_memory_status_get", + "parameters": [ + { + "in": "query", + "name": "key_id", + "required": false, + "schema": { + "anyOf": [ + { + "pattern": "^[a-f0-9]{64}$", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key Id" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MemoryStatus" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Status", + "tags": [ + "memory_v2" + ] + } + } + } + }, "policies": { "components": { "schemas": { diff --git a/litellm/proxy/memory/continuation.py b/litellm/proxy/memory/continuation.py index 96a679cf708..5d7ba0fdf1d 100644 --- a/litellm/proxy/memory/continuation.py +++ b/litellm/proxy/memory/continuation.py @@ -19,7 +19,9 @@ from litellm.repositories.unit_of_work import prisma_transaction _ITEMS: Final = TypeAdapter(tuple[object, ...]) _OBJECT: Final = TypeAdapter(dict[str, object]) _MAX_PATCH_BYTES: Final = 1024 * 1024 -_MAX_PATCHES: Final = 1000 +_MAX_PATCHES: Final = 256 +_MAX_NAMESPACE_BYTES: Final = 32 * 1024 * 1024 +_USAGE: Final = TypeAdapter(tuple[dict[str, int], ...]) async def cleanup_memory_continuations(prisma_client: object) -> None: @@ -204,31 +206,34 @@ class MemoryContinuations: key_id: Final = self.store.access.identity.key_id or self.store.access.identity.user_id or "" now: Final = datetime.now(timezone.utc) async with prisma_transaction(self.store.prisma_client) as transaction: - lock_key: Final = int(memory_digest("memory-continuation-quota", namespace, key_id)[:16], 16) - (1 << 63) + lock_key: Final = int(memory_digest("memory-continuation-quota", namespace)[:16], 16) - (1 << 63) await transaction.execute_raw("SELECT pg_advisory_xact_lock($1::bigint)", lock_key) table: Final = MemoryContinuationRepository(SimpleNamespace(db=transaction)).table await table.delete_many( where={ # mutable-ok: Prisma query and write JSON. "namespace": namespace, - "key_id": key_id, "expires_at": { # mutable-ok: Prisma query and write JSON. "lte": now }, } ) - count: Final = await table.count( - where={ # mutable-ok: Prisma query and write JSON. - "namespace": namespace, - "key_id": key_id, - "id": { # mutable-ok: Prisma query and write JSON. - "not_in": [ # mutable-ok: Prisma query and write JSON. - identifier for identifier, _ in payloads - ] - }, - } + usage: Final = _USAGE.validate_python( + await transaction.query_raw( + "SELECT COUNT(*) FILTER (WHERE key_id = $2)::int AS key_count, " + "COALESCE(SUM(octet_length(payload::text)), 0) + " + "(SELECT COALESCE(SUM(octet_length(value::text)), 0) " + "FROM jsonb_array_elements($4::jsonb)) AS bytes " + 'FROM "LiteLLM_MemoryContinuation" WHERE namespace = $1 AND NOT (id = ANY($3::text[]))', + namespace, + key_id, + [identifier for identifier, _ in payloads], # mutable-ok: Native Prisma array parameter. + "[" + ",".join(payload for _, payload in payloads) + "]", + ) ) - if count + len(payloads) > _MAX_PATCHES: + if usage[0]["key_count"] + len(payloads) > _MAX_PATCHES: raise HTTPException(status_code=429, detail="Too many active memory continuations for this key") + if usage[0]["bytes"] > _MAX_NAMESPACE_BYTES: + raise HTTPException(status_code=429, detail="Memory continuations exceed 32 megabytes for this scope") for identifier, payload in payloads: await table.upsert( where={ # mutable-ok: Prisma query and write JSON. diff --git a/litellm/proxy/memory/gateway.py b/litellm/proxy/memory/gateway.py index 443e101de26..09e85b16db8 100644 --- a/litellm/proxy/memory/gateway.py +++ b/litellm/proxy/memory/gateway.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.server_tools import ( append_server_reference, continue_server_tools, inject_server_tools, + trailing_system_messages, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.sse_keepalive import wrap_passthrough_sse_bytes_with_keepalive_pings @@ -68,6 +69,7 @@ class GatewayMemoryLoop: self.checkpoint = memory_digest(store.access.namespace, *prefix_hashes(self.visible_input, route)[-1:]) self.data: Mapping[str, object] = data self.baseline_length = 0 + self.replaced_input = 0 self.reflected = store.access.identity.read_only or ( data.get("tool_choice") not in (None, "auto") and object_value(data.get("tool_choice")).get("type") != "auto" @@ -115,7 +117,8 @@ class GatewayMemoryLoop: functions, MEMORY_READ_ONLY_WORKFLOW if self.store.access.identity.read_only else MEMORY_WORKFLOW, ) - self.baseline_length = len(transcript_items(injected, self.route)) + self.replaced_input = trailing_system_messages(injected, self.route) + self.baseline_length = len(transcript_items(injected, self.route)) - self.replaced_input catalog: Final = await memory_catalog(self.store, MemoryCatalogRequest(limit=12)) self.data = append_server_reference( injected, @@ -203,7 +206,7 @@ class GatewayMemoryLoop: visible: Final = response_messages(response, self.route) anchors: Final = prefix_hashes((*self.visible_input, *visible), self.route) patch: Final = MemoryContinuation( - replaces=len(visible), + replaces=len(visible) + self.replaced_input, replacement=transcript_items(self.data, self.route)[self.baseline_length :], upstream_ids=self.upstream_ids, pending_results=self.pending_results, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d3cb6839307..23de604a2e9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -582,7 +582,6 @@ from litellm.proxy.management_helpers.team_metadata_validation import ( TEAM_METADATA_VALIDATOR_REGISTRY, parse_team_metadata_schema, ) -from litellm.proxy.memory.management import router as memory_v2_router from litellm.proxy.memory.memory_endpoints import router as memory_router from litellm.proxy.middleware.billable_request_metrics_middleware import ( BillableRequestMetricsMiddleware, @@ -18803,7 +18802,6 @@ app.include_router(auto_router_management_router) app.include_router(tag_management_router) app.include_router(workflow_management_router) app.include_router(memory_router) -app.include_router(memory_v2_router) app.include_router(plugin_router) app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py index 2613ded4b91..f886d44cee7 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py +++ b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py @@ -53,6 +53,7 @@ def prisma_edge() -> MagicMock: table.count = AsyncMock(return_value=0) client.db.tx.return_value.__aenter__.return_value = client.db client.db.execute_raw = AsyncMock() + client.db.query_raw = AsyncMock(return_value=[{"key_count": 0, "bytes": 0}]) continuations = client.db.litellm_memorycontinuation continuations.find_many = AsyncMock(return_value=[]) continuations.find_first = AsyncMock(return_value=None) @@ -512,6 +513,72 @@ async def test_model_loop_is_bounded_and_search_results_reach_the_active_model(p prisma_edge.db.litellm_memorycontinuation.upsert.assert_not_awaited() +@pytest.mark.asyncio +async def test_trailing_system_messages_survive_client_tool_continuation(prisma_edge: MagicMock) -> None: + provider = FastAPI() + observed = [] + client_call = {"type": "tool_use", "id": "client_read", "name": "Read", "input": {"path": "README.md"}} + + @provider.post("/v1/messages") + async def model(incoming: Request): + body = await incoming.json() + observed.append(body) + messages = body["messages"] + assert all( + message["role"] != "system" or messages[index + 1]["role"] == "assistant" + for index, message in enumerate(messages[:-1]) + ) + return { + "id": "msg_" + str(len(observed)), + "role": "assistant", + "type": "message", + "stop_reason": "tool_use" if len(observed) == 1 else "end_turn", + "content": [client_call] if len(observed) == 1 else [{"type": "text", "text": "Read complete"}], + } + + prefix = { + "role": "user", + "content": [{"type": "text", "text": "Read README.md", "cache_control": {"type": "ephemeral"}}], + } + directive = {"role": "system", "content": "Use concise answers"} + original = { + "messages": [prefix, directive], + "tools": [{"name": "Read", "input_schema": {"type": "object"}}], + "tool_choice": {"type": "tool", "name": "Read"}, + } + first = GatewayMemoryLoop(provider, request(), original, "anthropic_messages", store(prisma_edge)) + async for _ in first.run(): + pass + saved = prisma_edge.db.litellm_memorycontinuation.upsert.call_args.kwargs + prisma_edge.db.litellm_memorycontinuation.find_many.return_value = [ + SimpleNamespace(id=saved["where"]["id"], payload=json.loads(saved["data"]["create"]["payload"])) + ] + following = { + **original, + "tool_choice": {"type": "none"}, + "messages": [ + prefix, + directive, + {"role": "assistant", "content": [client_call]}, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "client_read", "content": "File content"}], + }, + directive, + ], + } + second = GatewayMemoryLoop(provider, request(), following, "anthropic_messages", store(prisma_edge)) + async for _ in second.run(): + pass + assert len(observed) == 2 + assert observed[0]["messages"][0] == observed[1]["messages"][0] == prefix + assert observed[1]["messages"].count(directive) == 2 + assert observed[1]["messages"].count({"role": "assistant", "content": [client_call]}) == 1 + assert observed[1]["messages"][-1] == directive + assert observed[1]["messages"][-3]["content"][0]["tool_use_id"] == "client_read" + assert original["messages"] == [prefix, directive] + + @pytest.mark.asyncio @pytest.mark.parametrize("bad_id,count", [(True, 1), (False, 17)]) async def test_invalid_model_calls_are_rejected_before_storage( @@ -894,3 +961,38 @@ async def test_full_scope_blocks_creation_but_permits_correction_and_reclaimed_c table.create.return_value = row() assert (await store(prisma_edge).capture(_CAPTURE)).memory_id == "entry" table.create.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("key_count,used_bytes", [(256, 0), (0, 32 * 1024 * 1024 + 1)]) +async def test_continuation_quota_rejects_excess_without_writing( + prisma_edge: MagicMock, key_count: int, used_bytes: int +) -> None: + prisma_edge.db.query_raw.return_value = [{"key_count": key_count, "bytes": used_bytes}] + with pytest.raises(HTTPException) as exc: + await MemoryContinuations(store(prisma_edge), "aresponses").save("response", MemoryContinuation(replaces=1)) + assert exc.value.status_code == 429 + prisma_edge.db.litellm_memorycontinuation.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_continuation_quota_shares_namespace_lock_across_keys_and_allows_replacements( + prisma_edge: MagicMock, +) -> None: + user_policy = _POLICY.model_copy(update={"scope": "user"}) + prisma_edge.db.litellm_memorypolicy.find_many.return_value = [user_policy] + prisma_edge.db.query_raw.return_value = [{"key_count": 255, "bytes": 32 * 1024 * 1024}] + other_key = MemoryIdentity("b" * 64, "owner", "team", "project", "org", False) + for identity in (_IDENTITY, other_key): + continuations = MemoryContinuations( + MemoryStore(prisma_edge, MemoryAccess(identity, user_policy, False)), "aresponses" + ) + await continuations.save("response", MemoryContinuation(replaces=1, response={"text": "é漢字"})) + query = prisma_edge.db.query_raw.call_args.args + assert query[1:4] == (identity.namespace("user"), identity.key_id, [continuations.identifier("response")]) + assert json.loads(query[4])[0]["response"]["text"] == "é漢字" + locks = prisma_edge.db.execute_raw.call_args_list + assert locks[0] == locks[1] + cleanup = prisma_edge.db.litellm_memorycontinuation.delete_many.call_args.kwargs["where"] + assert cleanup["namespace"] == _IDENTITY.namespace("user") and "key_id" not in cleanup + assert prisma_edge.db.litellm_memorycontinuation.upsert.await_count == 2 diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_protocols.py b/tests/test_litellm/proxy/memory/test_memory_v2_protocols.py index d570dea3f19..b63b4ec1ffa 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_protocols.py +++ b/tests/test_litellm/proxy/memory/test_memory_v2_protocols.py @@ -91,6 +91,23 @@ def test_anthropic_continuation_preserves_signed_thinking_and_matches_tool_resul ] +@pytest.mark.parametrize("directive_only", [False, True]) +def test_memory_reference_keeps_anthropic_trailing_system_directives_valid(directive_only: bool) -> None: + prefix: Final = { + "role": "user", + "content": [{"type": "text", "text": "Read a file", "cache_control": {"type": "ephemeral"}}], + } + directive: Final = { + "role": "system", + "content": [] if directive_only else "Use concise answers", + "output_config": {"effort": "low"}, + } + original: Final = {"messages": [prefix, directive]} + result: Final = append_server_reference(original, "anthropic_messages", "Untrusted stored context") + assert result["messages"] == [prefix, {"role": "user", "content": "Untrusted stored context"}, directive] + assert original["messages"] == [prefix, directive] + + def test_responses_continuation_keeps_reasoning_and_function_call_output() -> None: output: Final = [ {"type": "reasoning", "id": "reason-1", "encrypted_content": "opaque-provider-data"},