hanzo-mcp 0.6.13__py3-none-any.whl → 0.7.1__py3-none-any.whl

hanzo_mcp/tools/common/batch_tool.py

@@ -13,6 +13,12 @@ from pydantic import Field
 
 from hanzo_mcp.tools.common.base import BaseTool
 from hanzo_mcp.tools.common.context import create_tool_context
+from hanzo_mcp.tools.common.truncate import truncate_response, estimate_tokens
+from hanzo_mcp.tools.common.fastmcp_pagination import (
+    create_paginated_response,
+    CursorData,
+    TokenAwarePaginator
+)
 
 
 class InvocationItem(TypedDict):
@@ -54,6 +60,14 @@ Invocations = Annotated[
     ),
 ]
 
+Cursor = Annotated[
+    str | None,
+    Field(
+        description="Pagination cursor to continue from previous batch results",
+        default=None,
+    ),
+]
+
 
 class BatchToolParams(TypedDict):
     """Parameters for the BatchTool.
@@ -61,10 +75,12 @@ class BatchToolParams(TypedDict):
     Attributes:
         description: A short (3-5 word) description of the batch operation
         invocations: The list of tool invocations to execute (required -- you MUST provide at least one tool invocation)
+        cursor: Optional pagination cursor
     """
 
     description: Description
     invocations: Invocations
+    cursor: Cursor
 
 
 @final
@@ -271,13 +287,78 @@ Not available: think,write,edit,multi_edit,notebook_edit
                     {"invocation": invocation, "result": f"Error: {error_message}"}
                 )
 
-        # Format the results
-        formatted_results = self._format_results(results)
-        await tool_ctx.info(
-            f"Batch operation '{description}' completed with {len(results)} results"
+        # Extract cursor if provided
+        cursor = params.get("cursor")
+        cursor_offset = 0
+        
+        # If cursor provided, we need to resume from where we left off
+        if cursor:
+            cursor_data = CursorData.from_cursor(cursor)
+            if cursor_data and cursor_data.offset < len(results):
+                # Skip already returned results
+                cursor_offset = cursor_data.offset
+                results = results[cursor_offset:]
+        
+        # Format results
+        formatted_results = []
+        for i, result in enumerate(results):
+            invocation = result["invocation"]
+            tool_name = invocation.get("tool_name", "unknown")
+            formatted_results.append({
+                "tool": tool_name,
+                "result": result["result"],
+                "index": i + cursor_offset
+            })
+        
+        # Create paginated response with token awareness
+        paginated_response = create_paginated_response(
+            formatted_results,
+            cursor=cursor,
+            use_token_limit=True
         )
-
-        return formatted_results
+        
+        # Convert paginated response to string format for MCP
+        if isinstance(paginated_response, dict) and "items" in paginated_response:
+            # Format the items as a readable string
+            result_parts = []
+            
+            # Add header
+            result_parts.append(f"=== Batch operation: {description} ===")
+            result_parts.append(f"Total invocations: {len(invocations)}")
+            result_parts.append(f"Showing results: {len(paginated_response['items'])} of {len(results)}")
+            if paginated_response.get('hasMore'):
+                result_parts.append(f"More results available - use cursor: {paginated_response.get('nextCursor')}")
+            result_parts.append("")
+            
+            # Format each result
+            for item in paginated_response['items']:
+                result_parts.append(f"### Result {item['index'] + 1}: {item['tool']}")
+                result_content = item['result']
+                
+                # Add the result content - use multi-line code blocks for code outputs
+                if isinstance(result_content, str) and "\n" in result_content:
+                    result_parts.append(f"```\n{result_content}\n```")
+                else:
+                    result_parts.append(str(result_content))
+                result_parts.append("")
+            
+            # Join all parts
+            formatted_output = "\n".join(result_parts)
+            
+            # If there's a next cursor, we need to preserve it in the response
+            # For now, append it as a note at the end
+            if paginated_response.get('hasMore') and paginated_response.get('nextCursor'):
+                formatted_output += f"\n\n[To continue, use cursor: {paginated_response['nextCursor']}]"
+            
+            await tool_ctx.info(
+                f"Batch operation '{description}' completed with {len(paginated_response['items'])} results"
+                f"{' (more available)' if paginated_response.get('hasMore') else ''}"
+            )
+            
+            return formatted_output
+        else:
+            # Fallback if pagination didn't work as expected
+            return self._format_results(results)
 
     def _format_results(self, results: list[dict[str, dict[str, Any]]]) -> str:
         """Format the results from multiple tool invocations.
@@ -295,11 +376,21 @@ Not available: think,write,edit,multi_edit,notebook_edit
 
             # Add the result header
             formatted_parts.append(f"### Result {i + 1}: {tool_name}")
+            
+            # Truncate individual results if they're too large
+            result_content = result["result"]
+            if len(result_content) > 50000:  # If individual result > 50k chars
+                result_content = truncate_response(
+                    result_content,
+                    max_tokens=5000,  # Limit individual results to ~5k tokens
+                    truncation_message=f"\n\n[Result from {tool_name} truncated. Use the tool directly with pagination/filtering for full output.]"
+                )
+            
             # Add the result content - use multi-line code blocks for code outputs
-            if "\n" in result["result"]:
-                formatted_parts.append(f"```\n{result['result']}\n```")
+            if "\n" in result_content:
+                formatted_parts.append(f"```\n{result_content}\n```")
             else:
-                formatted_parts.append(result["result"])
+                formatted_parts.append(result_content)
             # Add a separator
             formatted_parts.append("")
 
@@ -321,8 +412,9 @@ Not available: think,write,edit,multi_edit,notebook_edit
         async def batch(
             description: Description,
             invocations: Invocations,
+            cursor: Cursor,
             ctx: MCPContext
         ) -> str:
             return await tool_self.call(
-                ctx, description=description, invocations=invocations
+                ctx, description=description, invocations=invocations, cursor=cursor
             )

hanzo_mcp/tools/common/fastmcp_pagination.py

@@ -0,0 +1,369 @@
+"""FastMCP-compatible pagination implementation.
+
+This module provides pagination utilities optimized for FastMCP with minimal latency.
+"""
+
+import base64
+import json
+import time
+from typing import Any, Dict, List, Optional, Tuple, TypeVar, Generic, Union
+from dataclasses import dataclass, field
+from datetime import datetime
+import hashlib
+
+T = TypeVar('T')
+
+
+@dataclass
+class CursorData:
+    """Cursor data structure for efficient pagination."""
+    # Primary cursor fields (indexed)
+    last_id: Optional[str] = None
+    last_timestamp: Optional[float] = None
+    offset: int = 0
+    
+    # Metadata for validation and optimization
+    page_size: int = 100
+    sort_field: str = "id"
+    sort_order: str = "asc"
+    
+    # Security and validation
+    created_at: float = field(default_factory=time.time)
+    expires_at: Optional[float] = None
+    checksum: Optional[str] = None
+    
+    def to_cursor(self) -> str:
+        """Convert to opaque cursor string."""
+        data = {
+            "id": self.last_id,
+            "ts": self.last_timestamp,
+            "o": self.offset,
+            "ps": self.page_size,
+            "sf": self.sort_field,
+            "so": self.sort_order,
+            "ca": self.created_at,
+        }
+        if self.expires_at:
+            data["ea"] = self.expires_at
+            
+        # Add checksum for integrity
+        data_str = json.dumps(data, sort_keys=True, separators=(',', ':'))
+        data["cs"] = hashlib.md5(data_str.encode()).hexdigest()[:8]
+        
+        # Encode as base64
+        final_str = json.dumps(data, separators=(',', ':'))
+        return base64.urlsafe_b64encode(final_str.encode()).decode().rstrip('=')
+    
+    @classmethod
+    def from_cursor(cls, cursor: str) -> Optional['CursorData']:
+        """Parse cursor string back to CursorData."""
+        try:
+            # Add padding if needed
+            padding = 4 - (len(cursor) % 4)
+            if padding != 4:
+                cursor += '=' * padding
+                
+            decoded = base64.urlsafe_b64decode(cursor.encode())
+            data = json.loads(decoded)
+            
+            # Validate checksum
+            checksum = data.pop("cs", None)
+            if checksum:
+                data_str = json.dumps({k: v for k, v in data.items() if k != "cs"}, 
+                                    sort_keys=True, separators=(',', ':'))
+                expected = hashlib.md5(data_str.encode()).hexdigest()[:8]
+                if checksum != expected:
+                    return None
+            
+            # Check expiration
+            expires_at = data.get("ea")
+            if expires_at and time.time() > expires_at:
+                return None
+            
+            return cls(
+                last_id=data.get("id"),
+                last_timestamp=data.get("ts"),
+                offset=data.get("o", 0),
+                page_size=data.get("ps", 100),
+                sort_field=data.get("sf", "id"),
+                sort_order=data.get("so", "asc"),
+                created_at=data.get("ca", time.time()),
+                expires_at=expires_at
+            )
+        except Exception:
+            return None
+
+
+class FastMCPPaginator(Generic[T]):
+    """High-performance paginator for FastMCP responses."""
+    
+    def __init__(
+        self,
+        page_size: int = 100,
+        max_page_size: int = 1000,
+        cursor_ttl: int = 3600,  # 1 hour
+        enable_prefetch: bool = False
+    ):
+        """Initialize the paginator.
+        
+        Args:
+            page_size: Default page size
+            max_page_size: Maximum allowed page size
+            cursor_ttl: Cursor time-to-live in seconds
+            enable_prefetch: Enable prefetching for next page
+        """
+        self.page_size = page_size
+        self.max_page_size = max_page_size
+        self.cursor_ttl = cursor_ttl
+        self.enable_prefetch = enable_prefetch
+        self._cache: Dict[str, Any] = {}
+    
+    def paginate_list(
+        self,
+        items: List[T],
+        cursor: Optional[str] = None,
+        page_size: Optional[int] = None,
+        sort_key: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """Paginate a list with optimal performance.
+        
+        Args:
+            items: List to paginate
+            cursor: Optional cursor from previous request
+            page_size: Override default page size
+            sort_key: Sort field for consistent ordering
+            
+        Returns:
+            Dict with items and optional nextCursor
+        """
+        # Parse cursor or create new
+        cursor_data = CursorData.from_cursor(cursor) if cursor else CursorData()
+        
+        # Use provided page size or default
+        actual_page_size = min(
+            page_size or cursor_data.page_size or self.page_size,
+            self.max_page_size
+        )
+        
+        # Get starting position
+        start_idx = cursor_data.offset
+        
+        # Validate bounds
+        if start_idx >= len(items):
+            return {"items": [], "hasMore": False}
+        
+        # Slice the page
+        end_idx = min(start_idx + actual_page_size, len(items))
+        page_items = items[start_idx:end_idx]
+        
+        # Build response
+        response = {
+            "items": page_items,
+            "pageInfo": {
+                "startIndex": start_idx,
+                "endIndex": end_idx,
+                "pageSize": len(page_items),
+                "totalItems": len(items)
+            }
+        }
+        
+        # Create next cursor if more items exist
+        if end_idx < len(items):
+            next_cursor_data = CursorData(
+                offset=end_idx,
+                page_size=actual_page_size,
+                expires_at=time.time() + self.cursor_ttl if self.cursor_ttl else None
+            )
+            response["nextCursor"] = next_cursor_data.to_cursor()
+            response["hasMore"] = True
+        else:
+            response["hasMore"] = False
+        
+        return response
+    
+    def paginate_query(
+        self,
+        query_func,
+        cursor: Optional[str] = None,
+        page_size: Optional[int] = None,
+        **query_params
+    ) -> Dict[str, Any]:
+        """Paginate results from a query function.
+        
+        This is optimized for database queries using indexed fields.
+        
+        Args:
+            query_func: Function that accepts (last_id, last_timestamp, limit, **params)
+            cursor: Optional cursor
+            page_size: Override page size
+            **query_params: Additional query parameters
+            
+        Returns:
+            Paginated response
+        """
+        # Parse cursor
+        cursor_data = CursorData.from_cursor(cursor) if cursor else CursorData()
+        
+        # Determine page size
+        limit = min(
+            page_size or cursor_data.page_size or self.page_size,
+            self.max_page_size
+        )
+        
+        # Execute query with cursor position
+        results = query_func(
+            last_id=cursor_data.last_id,
+            last_timestamp=cursor_data.last_timestamp,
+            limit=limit + 1,  # Fetch one extra to detect more
+            **query_params
+        )
+        
+        # Check if there are more results
+        has_more = len(results) > limit
+        if has_more:
+            results = results[:limit]  # Remove the extra item
+        
+        # Build response
+        response = {
+            "items": results,
+            "pageInfo": {
+                "pageSize": len(results),
+                "hasMore": has_more
+            }
+        }
+        
+        # Create next cursor if needed
+        if has_more and results:
+            last_item = results[-1]
+            next_cursor_data = CursorData(
+                last_id=getattr(last_item, 'id', None),
+                last_timestamp=getattr(last_item, 'timestamp', None),
+                page_size=limit,
+                sort_field=cursor_data.sort_field,
+                sort_order=cursor_data.sort_order,
+                expires_at=time.time() + self.cursor_ttl if self.cursor_ttl else None
+            )
+            response["nextCursor"] = next_cursor_data.to_cursor()
+        
+        return response
+
+
+class TokenAwarePaginator:
+    """Paginator that respects token limits for LLM responses."""
+    
+    def __init__(self, max_tokens: int = 20000):
+        """Initialize token-aware paginator.
+        
+        Args:
+            max_tokens: Maximum tokens per response
+        """
+        self.max_tokens = max_tokens
+        self.paginator = FastMCPPaginator()
+    
+    def paginate_by_tokens(
+        self,
+        items: List[Any],
+        cursor: Optional[str] = None,
+        estimate_func=None
+    ) -> Dict[str, Any]:
+        """Paginate items based on token count.
+        
+        Args:
+            items: Items to paginate
+            cursor: Optional cursor
+            estimate_func: Function to estimate tokens for an item
+            
+        Returns:
+            Paginated response
+        """
+        from hanzo_mcp.tools.common.truncate import estimate_tokens
+        
+        # Default token estimation
+        if not estimate_func:
+            estimate_func = lambda x: estimate_tokens(json.dumps(x) if not isinstance(x, str) else x)
+        
+        # Parse cursor
+        cursor_data = CursorData.from_cursor(cursor) if cursor else CursorData()
+        start_idx = cursor_data.offset
+        
+        # Build page respecting token limit
+        page_items = []
+        current_tokens = 100  # Base overhead
+        current_idx = start_idx
+        
+        while current_idx < len(items) and current_tokens < self.max_tokens:
+            item = items[current_idx]
+            item_tokens = estimate_func(item)
+            
+            # Check if adding this item would exceed limit
+            if current_tokens + item_tokens > self.max_tokens and page_items:
+                break
+            
+            page_items.append(item)
+            current_tokens += item_tokens
+            current_idx += 1
+        
+        # Build response
+        response = {
+            "items": page_items,
+            "pageInfo": {
+                "itemCount": len(page_items),
+                "estimatedTokens": current_tokens,
+                "hasMore": current_idx < len(items)
+            }
+        }
+        
+        # Add next cursor if needed
+        if current_idx < len(items):
+            next_cursor_data = CursorData(offset=current_idx)
+            response["nextCursor"] = next_cursor_data.to_cursor()
+        
+        return response
+
+
+# FastMCP integration helpers
+def create_paginated_response(
+    items: Union[List[Any], Dict[str, Any], str],
+    cursor: Optional[str] = None,
+    page_size: int = 100,
+    use_token_limit: bool = True
+) -> Dict[str, Any]:
+    """Create a paginated response compatible with FastMCP.
+    
+    Args:
+        items: The items to paginate
+        cursor: Optional cursor from request
+        page_size: Items per page
+        use_token_limit: Whether to use token-based pagination
+        
+    Returns:
+        FastMCP-compatible paginated response
+    """
+    if use_token_limit:
+        paginator = TokenAwarePaginator()
+        
+        # Convert different types to list
+        if isinstance(items, str):
+            # Split string by lines for pagination
+            items = items.split('\n')
+        elif isinstance(items, dict):
+            # Convert dict to list of key-value pairs
+            items = [{"key": k, "value": v} for k, v in items.items()]
+        
+        return paginator.paginate_by_tokens(items, cursor)
+    else:
+        paginator = FastMCPPaginator(page_size=page_size)
+        
+        # Handle different input types
+        if isinstance(items, list):
+            return paginator.paginate_list(items, cursor, page_size)
+        else:
+            # Convert to list first
+            if isinstance(items, str):
+                items = items.split('\n')
+            elif isinstance(items, dict):
+                items = [{"key": k, "value": v} for k, v in items.items()]
+            else:
+                items = [items]
+            
+            return paginator.paginate_list(items, cursor, page_size)