Spaces:
Running
Running
File size: 2,826 Bytes
89e57b7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """Admission leases for HTTP and MCP content mutations (reads stay available)."""
import json
from starlette.responses import JSONResponse
from src.services.auth import request_permission
from .maintenance import admission
from .git import BackupError
READ_TOOLS = {'memory_search', 'memory_context', 'memory_neighbors', 'memory_entity',
'memory_task_status', 'search_memory', 'traverse_graph', 'get_brains',
'get_entity_info', 'get_instructions'}
READ_TOOLS.update({'list_brains', 'search_semantically', 'get_search_operation_instructions'})
class ContentAdmissionMiddleware:
def __init__(self, app):
self.app = app
async def __call__(self, scope, receive, send):
if scope['type'] != 'http':
return await self.app(scope, receive, send)
path, method = scope.get('path', ''), scope.get('method', 'GET')
write = request_permission(method, path) == 'memory:write'
if path.startswith(('/auth/', '/system/backup/', '/system/access-tokens')):
write = False
if '/mcp' in path and method == 'POST':
chunks = []
while True:
msg = await receive()
if msg['type'] != 'http.request':
return
chunks.append(msg.get('body', b''))
if sum(map(len, chunks)) > 2 * 1024 * 1024:
return await JSONResponse({'error': {'code': 'INVALID_REQUEST'}}, status_code=413)(scope, receive, send)
if not msg.get('more_body', False):
break
body = b''.join(chunks)
try:
payload = json.loads(body)
calls = payload if isinstance(payload, list) else [payload]
write = any(isinstance(c, dict) and c.get('method') == 'tools/call'
and (c.get('params') or {}).get('name') not in READ_TOOLS for c in calls)
except (ValueError, TypeError, AttributeError):
write = True
previous_receive = receive
delivered = False
async def replay():
nonlocal delivered
if not delivered:
delivered = True
return {'type': 'http.request', 'body': body, 'more_body': False}
return await previous_receive()
receive = replay
if not write:
return await self.app(scope, receive, send)
lease = admission()
try:
lease.__enter__()
except BackupError as exc:
return await JSONResponse({'detail': str(exc), 'error': {'code': str(exc)}}, status_code=503)(scope, receive, send)
try:
await self.app(scope, receive, send)
finally:
lease.__exit__(None, None, None)
|