hanzo-mcp 0.6.12__py3-none-any.whl → 0.7.0__py3-none-any.whl

hanzo_mcp/tools/common/paginated_response.py

@@ -0,0 +1,307 @@
+"""Automatic pagination response wrapper for MCP tools.
+
+This module provides utilities to automatically paginate tool responses
+when they exceed token limits.
+"""
+
+import json
+from typing import Any, Dict, List, Optional, Union
+from hanzo_mcp.tools.common.truncate import estimate_tokens
+from hanzo_mcp.tools.common.pagination import CursorManager
+
+
+class AutoPaginatedResponse:
+    """Automatically paginate responses that exceed token limits."""
+    
+    # MCP token limit with safety buffer
+    MAX_TOKENS = 20000  # Leave 5k buffer from 25k limit
+    
+    @staticmethod
+    def create_response(
+        content: Union[str, Dict[str, Any], List[Any]],
+        cursor: Optional[str] = None,
+        max_tokens: int = MAX_TOKENS
+    ) -> Dict[str, Any]:
+        """Create a response that automatically paginates if too large.
+        
+        Args:
+            content: The content to return
+            cursor: Optional cursor from request
+            max_tokens: Maximum tokens allowed in response
+            
+        Returns:
+            Dict with content and optional nextCursor
+        """
+        # Handle different content types
+        if isinstance(content, str):
+            return AutoPaginatedResponse._handle_string_response(
+                content, cursor, max_tokens
+            )
+        elif isinstance(content, list):
+            return AutoPaginatedResponse._handle_list_response(
+                content, cursor, max_tokens
+            )
+        elif isinstance(content, dict):
+            # If dict already has pagination info, return as-is
+            if "nextCursor" in content or "cursor" in content:
+                return content
+            # Otherwise treat as single item
+            return AutoPaginatedResponse._handle_dict_response(
+                content, cursor, max_tokens
+            )
+        else:
+            # Convert to string for other types
+            return AutoPaginatedResponse._handle_string_response(
+                str(content), cursor, max_tokens
+            )
+    
+    @staticmethod
+    def _handle_string_response(
+        content: str,
+        cursor: Optional[str],
+        max_tokens: int
+    ) -> Dict[str, Any]:
+        """Handle pagination for string responses."""
+        # Parse cursor to get offset
+        offset = 0
+        if cursor:
+            cursor_data = CursorManager.parse_cursor(cursor)
+            if cursor_data and "offset" in cursor_data:
+                offset = cursor_data["offset"]
+        
+        # For strings, paginate by lines
+        lines = content.split('\n')
+        
+        if offset >= len(lines):
+            return {"content": "", "message": "No more content"}
+        
+        # Build response line by line, checking tokens
+        result_lines = []
+        current_tokens = 0
+        line_index = offset
+        
+        # Add header if this is a continuation
+        if offset > 0:
+            header = f"[Continued from line {offset + 1}]\n"
+            current_tokens = estimate_tokens(header)
+            result_lines.append(header)
+        
+        while line_index < len(lines):
+            line = lines[line_index]
+            line_tokens = estimate_tokens(line + "\n")
+            
+            # Check if adding this line would exceed limit
+            if current_tokens + line_tokens > max_tokens:
+                # Need to paginate
+                if not result_lines:
+                    # Single line too long, truncate it
+                    truncated_line = line[:1000] + "... [line truncated]"
+                    result_lines.append(truncated_line)
+                    line_index += 1
+                break
+            
+            result_lines.append(line)
+            current_tokens += line_tokens
+            line_index += 1
+        
+        # Build response
+        response = {
+            "content": '\n'.join(result_lines)
+        }
+        
+        # Add pagination info
+        if line_index < len(lines):
+            response["nextCursor"] = CursorManager.create_offset_cursor(line_index)
+            response["pagination_info"] = {
+                "current_lines": f"{offset + 1}-{line_index}",
+                "total_lines": len(lines),
+                "has_more": True
+            }
+        else:
+            response["pagination_info"] = {
+                "current_lines": f"{offset + 1}-{len(lines)}",
+                "total_lines": len(lines),
+                "has_more": False
+            }
+        
+        return response
+    
+    @staticmethod
+    def _handle_list_response(
+        items: List[Any],
+        cursor: Optional[str],
+        max_tokens: int
+    ) -> Dict[str, Any]:
+        """Handle pagination for list responses."""
+        # Parse cursor to get offset
+        offset = 0
+        if cursor:
+            cursor_data = CursorManager.parse_cursor(cursor)
+            if cursor_data and "offset" in cursor_data:
+                offset = cursor_data["offset"]
+        
+        if offset >= len(items):
+            return {"items": [], "message": "No more items"}
+        
+        # Build response item by item, checking tokens
+        result_items = []
+        current_tokens = 100  # Base overhead
+        item_index = offset
+        
+        # Add header if continuation
+        header_obj = {}
+        if offset > 0:
+            header_obj["continuation_from"] = offset
+            current_tokens += 50
+        
+        while item_index < len(items):
+            item = items[item_index]
+            
+            # Estimate tokens for this item
+            item_str = json.dumps(item) if not isinstance(item, str) else item
+            item_tokens = estimate_tokens(item_str)
+            
+            # Check if adding this item would exceed limit
+            if current_tokens + item_tokens > max_tokens:
+                if not result_items:
+                    # Single item too large, truncate it
+                    if isinstance(item, str):
+                        truncated = item[:5000] + "... [truncated]"
+                        result_items.append(truncated)
+                    else:
+                        result_items.append({"error": "Item too large", "index": item_index})
+                    item_index += 1
+                break
+            
+            result_items.append(item)
+            current_tokens += item_tokens
+            item_index += 1
+        
+        # Build response
+        response = {
+            "items": result_items
+        }
+        
+        if header_obj:
+            response.update(header_obj)
+        
+        # Add pagination info
+        if item_index < len(items):
+            response["nextCursor"] = CursorManager.create_offset_cursor(item_index)
+            response["pagination_info"] = {
+                "returned_items": len(result_items),
+                "total_items": len(items),
+                "has_more": True,
+                "next_index": item_index
+            }
+        else:
+            response["pagination_info"] = {
+                "returned_items": len(result_items),
+                "total_items": len(items),
+                "has_more": False
+            }
+        
+        return response
+    
+    @staticmethod
+    def _handle_dict_response(
+        content: Dict[str, Any],
+        cursor: Optional[str],
+        max_tokens: int
+    ) -> Dict[str, Any]:
+        """Handle pagination for dict responses."""
+        # For dicts, check if it's too large as-is
+        content_str = json.dumps(content, indent=2)
+        content_tokens = estimate_tokens(content_str)
+        
+        if content_tokens <= max_tokens:
+            # Fits within limit
+            return content
+        
+        # Too large - need to paginate
+        # Strategy: Convert to key-value pairs and paginate
+        items = list(content.items())
+        offset = 0
+        
+        if cursor:
+            cursor_data = CursorManager.parse_cursor(cursor)
+            if cursor_data and "offset" in cursor_data:
+                offset = cursor_data["offset"]
+        
+        if offset >= len(items):
+            return {"content": {}, "message": "No more content"}
+        
+        # Build paginated dict
+        result = {}
+        current_tokens = 100  # Base overhead
+        
+        for i in range(offset, len(items)):
+            key, value = items[i]
+            
+            # Estimate tokens for this entry
+            entry_str = json.dumps({key: value})
+            entry_tokens = estimate_tokens(entry_str)
+            
+            if current_tokens + entry_tokens > max_tokens:
+                if not result:
+                    # Single entry too large
+                    result[key] = "[Value too large - use specific key access]"
+                break
+            
+            result[key] = value
+            current_tokens += entry_tokens
+        
+        # Wrap in response
+        response = {
+            "content": result
+        }
+        
+        # Add pagination info
+        processed = offset + len(result)
+        if processed < len(items):
+            response["nextCursor"] = CursorManager.create_offset_cursor(processed)
+            response["pagination_info"] = {
+                "keys_returned": len(result),
+                "total_keys": len(items),
+                "has_more": True
+            }
+        else:
+            response["pagination_info"] = {
+                "keys_returned": len(result),
+                "total_keys": len(items),
+                "has_more": False
+            }
+        
+        return response
+
+
+def paginate_if_needed(
+    response: Any,
+    cursor: Optional[str] = None,
+    force_pagination: bool = False
+) -> Union[str, Dict[str, Any]]:
+    """Wrap a response with automatic pagination if needed.
+    
+    Args:
+        response: The response to potentially paginate
+        cursor: Optional cursor from request
+        force_pagination: Force pagination even for small responses
+        
+    Returns:
+        Original response if small enough, otherwise paginated dict
+    """
+    # Quick check - if response is already paginated, return as-is
+    if isinstance(response, dict) and ("nextCursor" in response or "pagination_info" in response):
+        return response
+    
+    # For small responses, don't paginate unless forced
+    if not force_pagination:
+        try:
+            response_str = json.dumps(response) if not isinstance(response, str) else response
+            if len(response_str) < 10000:  # Quick heuristic
+                return response
+        except:
+            pass
+    
+    # Create paginated response
+    return AutoPaginatedResponse.create_response(response, cursor)

hanzo_mcp/tools/common/pagination.py

@@ -0,0 +1,226 @@
+"""Pagination utilities for MCP tools.
+
+This module provides utilities for implementing cursor-based pagination
+according to the MCP pagination protocol.
+"""
+
+import base64
+import json
+from typing import Any, Dict, List, Optional, Tuple, TypeVar, Generic
+from dataclasses import dataclass
+
+
+T = TypeVar('T')
+
+
+@dataclass
+class PaginationParams:
+    """Parameters for pagination."""
+    cursor: Optional[str] = None
+    page_size: int = 100  # Default page size
+
+
+@dataclass
+class PaginatedResponse(Generic[T]):
+    """A paginated response containing items and optional next cursor."""
+    items: List[T]
+    next_cursor: Optional[str] = None
+    
+    def to_dict(self, items_key: str = "items") -> Dict[str, Any]:
+        """Convert to dictionary format for MCP response.
+        
+        Args:
+            items_key: The key to use for items in the response
+            
+        Returns:
+            Dictionary with items and optional nextCursor
+        """
+        result = {items_key: self.items}
+        if self.next_cursor:
+            result["nextCursor"] = self.next_cursor
+        return result
+
+
+class CursorManager:
+    """Manages cursor creation and parsing for pagination."""
+    
+    @staticmethod
+    def create_cursor(data: Dict[str, Any]) -> str:
+        """Create an opaque cursor from data.
+        
+        Args:
+            data: Data to encode in the cursor
+            
+        Returns:
+            Base64-encoded cursor string
+        """
+        json_data = json.dumps(data, separators=(',', ':'))
+        return base64.b64encode(json_data.encode()).decode()
+    
+    @staticmethod
+    def parse_cursor(cursor: str) -> Optional[Dict[str, Any]]:
+        """Parse a cursor string back to data.
+        
+        Args:
+            cursor: Base64-encoded cursor string
+            
+        Returns:
+            Decoded data or None if invalid
+        """
+        try:
+            decoded = base64.b64decode(cursor.encode()).decode()
+            return json.loads(decoded)
+        except (ValueError, json.JSONDecodeError):
+            return None
+    
+    @staticmethod
+    def create_offset_cursor(offset: int) -> str:
+        """Create a cursor for offset-based pagination.
+        
+        Args:
+            offset: The offset for the next page
+            
+        Returns:
+            Cursor string
+        """
+        return CursorManager.create_cursor({"offset": offset})
+    
+    @staticmethod
+    def parse_offset_cursor(cursor: Optional[str]) -> int:
+        """Parse an offset cursor.
+        
+        Args:
+            cursor: Cursor string or None
+            
+        Returns:
+            Offset value (0 if cursor is None or invalid)
+        """
+        if not cursor:
+            return 0
+        
+        data = CursorManager.parse_cursor(cursor)
+        if not data or "offset" not in data:
+            return 0
+        
+        return int(data["offset"])
+
+
+class Paginator(Generic[T]):
+    """Generic paginator for any list of items."""
+    
+    def __init__(self, items: List[T], page_size: int = 100):
+        """Initialize paginator.
+        
+        Args:
+            items: List of items to paginate
+            page_size: Number of items per page
+        """
+        self.items = items
+        self.page_size = page_size
+    
+    def get_page(self, cursor: Optional[str] = None) -> PaginatedResponse[T]:
+        """Get a page of results.
+        
+        Args:
+            cursor: Optional cursor for the page
+            
+        Returns:
+            Paginated response with items and next cursor
+        """
+        offset = CursorManager.parse_offset_cursor(cursor)
+        
+        # Get the page of items
+        start = offset
+        end = min(start + self.page_size, len(self.items))
+        page_items = self.items[start:end]
+        
+        # Create next cursor if there are more items
+        next_cursor = None
+        if end < len(self.items):
+            next_cursor = CursorManager.create_offset_cursor(end)
+        
+        return PaginatedResponse(items=page_items, next_cursor=next_cursor)
+
+
+class StreamPaginator(Generic[T]):
+    """Paginator for streaming/generator-based results."""
+    
+    def __init__(self, page_size: int = 100):
+        """Initialize stream paginator.
+        
+        Args:
+            page_size: Number of items per page
+        """
+        self.page_size = page_size
+    
+    def paginate_stream(
+        self, 
+        stream_generator,
+        cursor: Optional[str] = None
+    ) -> PaginatedResponse[T]:
+        """Paginate results from a stream/generator.
+        
+        Args:
+            stream_generator: Generator function that yields items
+            cursor: Optional cursor for resuming
+            
+        Returns:
+            Paginated response
+        """
+        items = []
+        skip_count = 0
+        
+        # Parse cursor to get skip count
+        if cursor:
+            cursor_data = CursorManager.parse_cursor(cursor)
+            if cursor_data and "skip" in cursor_data:
+                skip_count = cursor_data["skip"]
+        
+        # Skip items based on cursor
+        item_count = 0
+        for item in stream_generator():
+            if item_count < skip_count:
+                item_count += 1
+                continue
+            
+            items.append(item)
+            if len(items) >= self.page_size:
+                # We have a full page, create cursor for next page
+                next_cursor = CursorManager.create_cursor({
+                    "skip": skip_count + len(items)
+                })
+                return PaginatedResponse(items=items, next_cursor=next_cursor)
+        
+        # No more items
+        return PaginatedResponse(items=items, next_cursor=None)
+
+
+def paginate_list(
+    items: List[T],
+    cursor: Optional[str] = None,
+    page_size: int = 100
+) -> PaginatedResponse[T]:
+    """Convenience function to paginate a list.
+    
+    Args:
+        items: List to paginate
+        cursor: Optional cursor
+        page_size: Items per page
+        
+    Returns:
+        Paginated response
+    """
+    paginator = Paginator(items, page_size)
+    return paginator.get_page(cursor)
+
+
+def validate_cursor(cursor: str) -> bool:
+    """Validate that a cursor is properly formatted.
+    
+    Args:
+        cursor: Cursor string to validate
+        
+    Returns:
+        True if valid, False otherwise
+    """
+    return CursorManager.parse_cursor(cursor) is not None

hanzo_mcp/tools/common/permissions.py

@@ -1,4 +1,4 @@
-"""Permission system for the Hanzo MCP server."""
+"""Permission system for the Hanzo AI server."""
 
 import json
 import os