marm-mcp

"""Response size limiting utilities for MCP compliance.""" import json import sys from typing import Dict, List, Any, Tuple class MCPResponseLimiter: """Ensures MCP responses stay under 1MB limit.""" # MCP response size limit (1MB = 1,048,576 bytes) MAX_RESPONSE_SIZE = 1048576 # Safe buffer to account for JSON overhead and metadata SAFE_BUFFER = 48576 # ~47KB buffer CONTENT_LIMIT = MAX_RESPONSE_SIZE - SAFE_BUFFER # ~1000KB for actual content @staticmethod def estimate_response_size(response_data: Dict[Any, Any]) -> int: """Estimate JSON response size in bytes.""" try: json_str = json.dumps(response_data, ensure_ascii=False) return len(json_str.encode('utf-8')) except Exception as e: # Fallback estimation if JSON serialization fails print(f"Warning: JSON serialization failed ({e}), using inaccurate fallback for size estimation") return sys.getsizeof(str(response_data)) @staticmethod def truncate_content(content: str, max_chars: int, preserve_start: bool = True) -> str: """Intelligently truncate content while preserving meaning.""" if len(content) <= max_chars: return content if preserve_start: # Keep beginning, add ellipsis return content[:max_chars-3] + "..." else: # Keep end, add ellipsis at start return "..." + content[-(max_chars-3):] @staticmethod def truncate_memory_content(memory: Dict[str, Any], max_content_chars: int, preserve_start: bool = True) -> Dict[str, Any]: """Truncate a single memory's content field.""" truncated = memory.copy() if 'content' in truncated: original_length = len(truncated['content']) # Choose truncation strategy based on content type # Error logs might benefit from preserving the end (final error message) content_type = memory.get('context_type', 'general') if content_type in ['error', 'debug', 'log'] and preserve_start is True: # For error content, preserve end by default unless explicitly requested otherwise preserve_start = False truncated['content'] = MCPResponseLimiter.truncate_content( truncated['content'], max_content_chars, preserve_start ) # Add truncation indicator if original_length > max_content_chars: truncated['_truncated'] = True truncated['_original_length'] = original_length truncated['_truncation_strategy'] = 'start' if preserve_start else 'end' return truncated @classmethod def limit_memory_response(cls, memories: List[Dict[str, Any]], response_metadata: Dict[str, Any], preserve_start: bool = True) -> Tuple[List[Dict[str, Any]], bool]: """ Limit memory list response to stay under MCP size limits. Returns (limited_memories, was_truncated) """ if not memories: return memories, False # Estimate base response size without content base_response = response_metadata.copy() base_response['results'] = [] base_size = cls.estimate_response_size(base_response) # Calculate available space for memory content available_space = cls.CONTENT_LIMIT - base_size # If base response is already too big, return minimal response if available_space <= 0: return [], True # Try to fit memories, truncating content as needed limited_memories = [] current_size = base_size was_truncated = False # Calculate rough max content per memory estimated_content_per_memory = available_space // len(memories) # Reserve space for JSON overhead (field names, quotes, commas, etc.) max_content_per_memory = max(100, estimated_content_per_memory - 200) for memory in memories: # Truncate this memory's content with smart strategy selection truncated_memory = cls.truncate_memory_content( memory, max_content_per_memory, preserve_start ) # Test if adding this memory exceeds limit test_memories = limited_memories + [truncated_memory] test_response = response_metadata.copy() test_response['results'] = test_memories test_size = cls.estimate_response_size(test_response) if test_size > cls.CONTENT_LIMIT: # Can't fit this memory, stop here was_truncated = True break # Memory fits, add it limited_memories.append(truncated_memory) current_size = test_size # Check if we truncated content if '_truncated' in truncated_memory: was_truncated = True return limited_memories, was_truncated @classmethod def limit_context_bridge_response(cls, related_content: List[Dict[str, Any]], response_metadata: Dict[str, Any]) -> Tuple[List[Dict[str, Any]], bool]: """Limit context bridge response with multiple content types.""" if not related_content: return related_content, False # Similar to memory response but handles mixed content types base_response = response_metadata.copy() base_response['related_content'] = [] base_size = cls.estimate_response_size(base_response) available_space = cls.CONTENT_LIMIT - base_size if available_space <= 0: return [], True limited_content = [] was_truncated = False max_content_per_item = max(150, (available_space // len(related_content)) - 300) for item in related_content: # Truncate content field in context items truncated_item = item.copy() if 'content' in truncated_item: original_length = len(truncated_item['content']) truncated_item['content'] = cls.truncate_content( truncated_item['content'], max_content_per_item ) if original_length > max_content_per_item: truncated_item['_truncated'] = True was_truncated = True # Test size test_content = limited_content + [truncated_item] test_response = response_metadata.copy() test_response['related_content'] = test_content test_size = cls.estimate_response_size(test_response) if test_size > cls.CONTENT_LIMIT: was_truncated = True break limited_content.append(truncated_item) return limited_content, was_truncated @classmethod def add_truncation_notice(cls, response: Dict[str, Any], was_truncated: bool, total_available: int = None) -> Dict[str, Any]: """Add truncation notice to response if content was limited.""" if was_truncated: response['_mcp_truncated'] = True response['_truncation_reason'] = "Response limited to 1MB for MCP compliance" if total_available: response['_total_available'] = total_available response['message'] = response.get('message', '') + f" (showing partial results due to size limit, {total_available} total available)" return response