claude-mpm 3.9.11__py3-none-any.whl → 4.0.4__py3-none-any.whl

claude_mpm/services/infrastructure/context_preservation.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 """Context Preservation service for handling Claude conversation data.
 
 This service specializes in parsing and preserving Claude's conversation
@@ -10,86 +12,115 @@ Design Principles:
 - Graceful handling of corrupted data
 """
 
-import json
-import logging
 import gzip
+import json
 import os
 import shutil
 import tempfile
 from datetime import datetime, timedelta
-from pathlib import Path
-from typing import Dict, Any, List, Optional, Iterator, Tuple
+from typing import Any, Dict, Iterator, List, Optional, Tuple
+
 import ijson  # For streaming JSON parsing
+from dataclasses import dataclass, field
 
 from claude_mpm.services.core.base import BaseService
-from claude_mpm.models.state_models import ConversationState, ConversationContext
+
+
+@dataclass
+class ConversationContext:
+    """Context information for a single conversation."""
+
+    conversation_id: str
+    title: str
+    message_count: int
+    last_message_time: float
+    file_references: List[str] = field(default_factory=list)
+    open_tabs: List[str] = field(default_factory=list)
+    tags: List[str] = field(default_factory=list)
+    is_active: bool = False
+
+
+@dataclass
+class ConversationState:
+    """Claude conversation state and context."""
+
+    active_conversation_id: Optional[str]
+    active_conversation: Optional[ConversationContext]
+    recent_conversations: List[ConversationContext]
+    total_conversations: int
+    total_storage_mb: float
+    preferences: Dict[str, Any]
+    open_files: List[str]
+    recent_files: List[str]
+    pinned_files: List[str]
 
 
 class ContextPreservationService(BaseService):
     """Service for preserving and managing Claude conversation context."""
-    
+
     def __init__(self, claude_dir: Optional[Path] = None):
         """Initialize Context Preservation service.
-        
+
         Args:
             claude_dir: Claude configuration directory (default: ~/.claude)
         """
         super().__init__("ContextPreservation")
-        
+
         # Claude configuration paths
         self.claude_dir = claude_dir or Path.home() / ".claude"
         self.claude_json_path = self.claude_dir / ".claude.json"
         self.claude_backup_dir = self.claude_dir / "backups"
-        
+
         # Size thresholds
         self.large_file_threshold_mb = 100
         self.compression_threshold_mb = 50
-        
+
         # Context extraction limits
         self.max_conversations_to_extract = 10
         self.max_messages_per_conversation = 100
         self.max_file_references = 1000
-        
+
         # Statistics
         self.files_processed = 0
         self.total_size_processed_mb = 0.0
-        
+
         self.log_info(f"Context Preservation initialized for: {self.claude_dir}")
-    
+
     async def initialize(self) -> bool:
         """Initialize the Context Preservation service.
-        
+
         Returns:
             True if initialization successful
         """
         try:
             self.log_info("Initializing Context Preservation service")
-            
+
             # Ensure backup directory exists
             self.claude_backup_dir.mkdir(parents=True, exist_ok=True)
-            
+
             # Check Claude configuration
             if self.claude_json_path.exists():
                 size_mb = self.claude_json_path.stat().st_size / (1024 * 1024)
                 self.log_info(f"Found Claude configuration: {size_mb:.2f}MB")
             else:
                 self.log_warning("Claude configuration not found")
-            
+
             self._initialized = True
             self.log_info("Context Preservation service initialized successfully")
             return True
-            
+
         except Exception as e:
             self.log_error(f"Failed to initialize Context Preservation: {e}")
             return False
-    
-    async def parse_claude_json(self, 
-                               extract_full: bool = False) -> Optional[ConversationState]:
+
+    async def parse_claude_json(
+        self, extract_full: bool = False
+    ) -> Optional[ConversationState]:
         """Parse Claude's .claude.json file safely.
-        
+
         Args:
             extract_full: Whether to extract full conversation data
-            
+
         Returns:
             ConversationState object or None if parsing failed
         """
@@ -97,223 +128,229 @@ class ContextPreservationService(BaseService):
             if not self.claude_json_path.exists():
                 self.log_debug("Claude configuration file not found")
                 return self._empty_conversation_state()
-            
+
             file_size_mb = self.claude_json_path.stat().st_size / (1024 * 1024)
             self.log_info(f"Parsing Claude configuration: {file_size_mb:.2f}MB")
-            
+
             # Choose parsing strategy based on file size
             if file_size_mb > self.large_file_threshold_mb:
                 self.log_info("Using streaming parser for large file")
                 return await self._parse_large_claude_json(extract_full)
             else:
                 return await self._parse_standard_claude_json(extract_full)
-                
+
         except Exception as e:
             self.log_error(f"Failed to parse Claude JSON: {e}")
             return None
-    
+
     async def extract_active_conversation(self) -> Optional[ConversationContext]:
         """Extract only the active conversation context.
-        
+
         Returns:
             ConversationContext for active conversation or None
         """
         try:
             if not self.claude_json_path.exists():
                 return None
-            
+
             # Use streaming to find active conversation
-            with open(self.claude_json_path, 'rb') as f:
+            with open(self.claude_json_path, "rb") as f:
                 parser = ijson.parse(f)
-                
+
                 active_conv_id = None
                 in_conversations = False
                 current_conv = {}
-                
+
                 for prefix, event, value in parser:
                     # Get active conversation ID
-                    if prefix == 'activeConversationId':
+                    if prefix == "activeConversationId":
                         active_conv_id = value
-                    
+
                     # Parse conversations array
-                    elif prefix.startswith('conversations.item'):
-                        if event == 'map_key':
+                    elif prefix.startswith("conversations.item"):
+                        if event == "map_key":
                             current_key = value
-                        elif active_conv_id and current_conv.get('id') == active_conv_id:
+                        elif (
+                            active_conv_id and current_conv.get("id") == active_conv_id
+                        ):
                             # Found active conversation
                             return self._create_conversation_context(current_conv)
-                
+
                 return None
-                
+
         except Exception as e:
             self.log_error(f"Failed to extract active conversation: {e}")
             return None
-    
-    async def compress_conversation_history(self,
-                                           keep_recent_days: int = 7) -> bool:
+
+    async def compress_conversation_history(self, keep_recent_days: int = 7) -> bool:
         """Compress large conversation histories.
-        
+
         Args:
             keep_recent_days: Days of recent conversations to keep uncompressed
-            
+
         Returns:
             True if compression successful
         """
         try:
             if not self.claude_json_path.exists():
                 return False
-            
+
             file_size_mb = self.claude_json_path.stat().st_size / (1024 * 1024)
-            
+
             if file_size_mb < self.compression_threshold_mb:
                 self.log_debug(f"File too small for compression: {file_size_mb:.2f}MB")
                 return False
-            
+
             self.log_info(f"Compressing conversation history: {file_size_mb:.2f}MB")
-            
+
             # Create backup first
             backup_path = await self._create_backup()
-            
+
             # Load and filter conversations
             cutoff_time = datetime.now().timestamp() - (keep_recent_days * 86400)
-            
-            with open(self.claude_json_path, 'r') as f:
+
+            with open(self.claude_json_path, "r") as f:
                 data = json.load(f)
-            
-            original_count = len(data.get('conversations', []))
-            
+
+            original_count = len(data.get("conversations", []))
+
             # Filter conversations
             recent_conversations = []
             archived_conversations = []
-            
-            for conv in data.get('conversations', []):
-                updated_at = conv.get('updatedAt', 0) / 1000  # Convert from ms
+
+            for conv in data.get("conversations", []):
+                updated_at = conv.get("updatedAt", 0) / 1000  # Convert from ms
                 if updated_at >= cutoff_time:
                     recent_conversations.append(conv)
                 else:
                     # Create minimal version for archive
-                    archived_conversations.append({
-                        'id': conv.get('id'),
-                        'title': conv.get('title'),
-                        'createdAt': conv.get('createdAt'),
-                        'updatedAt': conv.get('updatedAt'),
-                        'messageCount': len(conv.get('messages', [])),
-                        'archived': True
-                    })
-            
+                    archived_conversations.append(
+                        {
+                            "id": conv.get("id"),
+                            "title": conv.get("title"),
+                            "createdAt": conv.get("createdAt"),
+                            "updatedAt": conv.get("updatedAt"),
+                            "messageCount": len(conv.get("messages", [])),
+                            "archived": True,
+                        }
+                    )
+
             # Update data with filtered conversations
-            data['conversations'] = recent_conversations
-            data['archivedConversations'] = archived_conversations
-            
+            data["conversations"] = recent_conversations
+            data["archivedConversations"] = archived_conversations
+
             # Write compressed version
-            temp_path = self.claude_json_path.with_suffix('.tmp')
-            with open(temp_path, 'w') as f:
-                json.dump(data, f, separators=(',', ':'))  # Compact format
-            
+            temp_path = self.claude_json_path.with_suffix(".tmp")
+            with open(temp_path, "w") as f:
+                json.dump(data, f, separators=(",", ":"))  # Compact format
+
             # Replace original
             temp_path.replace(self.claude_json_path)
-            
+
             new_size_mb = self.claude_json_path.stat().st_size / (1024 * 1024)
             reduction_pct = ((file_size_mb - new_size_mb) / file_size_mb) * 100
-            
-            self.log_info(f"Compression complete: {original_count} -> {len(recent_conversations)} "
-                         f"conversations, {file_size_mb:.2f}MB -> {new_size_mb:.2f}MB "
-                         f"({reduction_pct:.1f}% reduction)")
-            
+
+            self.log_info(
+                f"Compression complete: {original_count} -> {len(recent_conversations)} "
+                f"conversations, {file_size_mb:.2f}MB -> {new_size_mb:.2f}MB "
+                f"({reduction_pct:.1f}% reduction)"
+            )
+
             return True
-            
+
         except Exception as e:
             self.log_error(f"Failed to compress conversation history: {e}")
             return False
-    
-    async def handle_file_references(self, 
-                                    conversation: Dict[str, Any]) -> List[str]:
+
+    async def handle_file_references(self, conversation: Dict[str, Any]) -> List[str]:
         """Extract and validate file references from conversation.
-        
+
         Args:
             conversation: Conversation data dictionary
-            
+
         Returns:
             List of valid file paths referenced in conversation
         """
         try:
             files = set()
-            
+
             # Extract from messages
-            for message in conversation.get('messages', [])[:self.max_messages_per_conversation]:
+            for message in conversation.get("messages", [])[
+                : self.max_messages_per_conversation
+            ]:
                 # Extract from content
-                content = message.get('content', '')
+                content = message.get("content", "")
                 if isinstance(content, str):
                     files.update(self._extract_file_paths(content))
-                
+
                 # Extract from attachments
-                for attachment in message.get('attachments', []):
-                    if attachment.get('type') == 'file':
-                        file_path = attachment.get('path')
+                for attachment in message.get("attachments", []):
+                    if attachment.get("type") == "file":
+                        file_path = attachment.get("path")
                         if file_path:
                             files.add(file_path)
-            
+
             # Validate file existence
             valid_files = []
-            for file_path in list(files)[:self.max_file_references]:
+            for file_path in list(files)[: self.max_file_references]:
                 try:
                     if Path(file_path).exists():
                         valid_files.append(file_path)
                 except:
                     pass  # Invalid path
-            
+
             return valid_files
-            
+
         except Exception as e:
             self.log_error(f"Failed to handle file references: {e}")
             return []
-    
+
     async def preserve_user_preferences(self) -> Dict[str, Any]:
         """Extract and preserve user preferences and settings.
-        
+
         Returns:
             Dictionary of user preferences
         """
         try:
             if not self.claude_json_path.exists():
                 return {}
-            
+
             # Use streaming to extract preferences
             preferences = {}
-            
-            with open(self.claude_json_path, 'rb') as f:
+
+            with open(self.claude_json_path, "rb") as f:
                 parser = ijson.parse(f)
-                
+
                 for prefix, event, value in parser:
-                    if prefix.startswith('preferences'):
+                    if prefix.startswith("preferences"):
                         # Extract preference key and value
-                        if event == 'map_key':
+                        if event == "map_key":
                             current_key = value
-                        elif event in ('string', 'number', 'boolean'):
+                        elif event in ("string", "number", "boolean"):
                             preferences[current_key] = value
-            
+
             self.log_debug(f"Preserved {len(preferences)} user preferences")
             return preferences
-            
+
         except Exception as e:
             self.log_error(f"Failed to preserve user preferences: {e}")
             return {}
-    
-    async def _parse_standard_claude_json(self, 
-                                         extract_full: bool) -> ConversationState:
+
+    async def _parse_standard_claude_json(
+        self, extract_full: bool
+    ) -> ConversationState:
         """Parse Claude JSON using standard JSON parser."""
         try:
-            with open(self.claude_json_path, 'r') as f:
+            with open(self.claude_json_path, "r") as f:
                 data = json.load(f)
-            
+
             return await self._extract_conversation_state(data, extract_full)
-            
+
         except json.JSONDecodeError as e:
             self.log_error(f"JSON decode error: {e}")
             return self._empty_conversation_state()
-    
-    async def _parse_large_claude_json(self, 
-                                      extract_full: bool) -> ConversationState:
+
+    async def _parse_large_claude_json(self, extract_full: bool) -> ConversationState:
         """Parse large Claude JSON using streaming parser."""
         try:
             # Extract key data using streaming
@@ -321,56 +358,56 @@ class ContextPreservationService(BaseService):
             conversations = []
             preferences = {}
             open_files = []
-            
-            with open(self.claude_json_path, 'rb') as f:
+
+            with open(self.claude_json_path, "rb") as f:
                 parser = ijson.parse(f)
-                
+
                 conversation_count = 0
                 current_conv = {}
                 in_conversation = False
-                
+
                 for prefix, event, value in parser:
                     # Limit conversations extracted
                     if conversation_count >= self.max_conversations_to_extract:
                         break
-                    
+
                     # Extract active conversation ID
-                    if prefix == 'activeConversationId':
+                    if prefix == "activeConversationId":
                         active_conv_id = value
-                    
+
                     # Extract conversations
-                    elif prefix.startswith('conversations.item'):
-                        if event == 'start_map':
+                    elif prefix.startswith("conversations.item"):
+                        if event == "start_map":
                             in_conversation = True
                             current_conv = {}
-                        elif event == 'end_map':
+                        elif event == "end_map":
                             if in_conversation and current_conv:
                                 conversations.append(current_conv)
                                 conversation_count += 1
                             in_conversation = False
                             current_conv = {}
-                        elif in_conversation and event == 'map_key':
+                        elif in_conversation and event == "map_key":
                             current_key = value
                         elif in_conversation and current_key:
                             current_conv[current_key] = value
-                    
+
                     # Extract open files
-                    elif prefix.startswith('openFiles.item'):
-                        if event == 'string':
+                    elif prefix.startswith("openFiles.item"):
+                        if event == "string":
                             open_files.append(value)
-            
+
             # Find active conversation
             active_conv = None
             recent_convs = []
-            
+
             for conv in conversations:
-                if conv.get('id') == active_conv_id:
+                if conv.get("id") == active_conv_id:
                     active_conv = self._create_conversation_context(conv)
                 else:
                     recent_convs.append(self._create_conversation_context(conv))
-            
+
             file_size_mb = self.claude_json_path.stat().st_size / (1024 * 1024)
-            
+
             return ConversationState(
                 active_conversation_id=active_conv_id,
                 active_conversation=active_conv,
@@ -380,77 +417,74 @@ class ContextPreservationService(BaseService):
                 preferences=preferences,
                 open_files=open_files[:100],
                 recent_files=[],
-                pinned_files=[]
+                pinned_files=[],
             )
-            
+
         except Exception as e:
             self.log_error(f"Failed to parse large Claude JSON: {e}")
             return self._empty_conversation_state()
-    
-    async def _extract_conversation_state(self, 
-                                         data: Dict[str, Any],
-                                         extract_full: bool) -> ConversationState:
+
+    async def _extract_conversation_state(
+        self, data: Dict[str, Any], extract_full: bool
+    ) -> ConversationState:
         """Extract conversation state from parsed data."""
         try:
-            conversations = data.get('conversations', [])
-            active_conv_id = data.get('activeConversationId')
-            
+            conversations = data.get("conversations", [])
+            active_conv_id = data.get("activeConversationId")
+
             # Find active conversation
             active_conv = None
             if active_conv_id:
                 for conv in conversations:
-                    if conv.get('id') == active_conv_id:
+                    if conv.get("id") == active_conv_id:
                         active_conv = self._create_conversation_context(conv)
                         break
-            
+
             # Get recent conversations
             recent_convs = []
             if extract_full:
                 sorted_convs = sorted(
-                    conversations,
-                    key=lambda c: c.get('updatedAt', 0),
-                    reverse=True
-                )[:self.max_conversations_to_extract]
-                
+                    conversations, key=lambda c: c.get("updatedAt", 0), reverse=True
+                )[: self.max_conversations_to_extract]
+
                 for conv in sorted_convs:
-                    if conv.get('id') != active_conv_id:
+                    if conv.get("id") != active_conv_id:
                         recent_convs.append(self._create_conversation_context(conv))
-            
+
             file_size_mb = self.claude_json_path.stat().st_size / (1024 * 1024)
-            
+
             return ConversationState(
                 active_conversation_id=active_conv_id,
                 active_conversation=active_conv,
                 recent_conversations=recent_convs,
                 total_conversations=len(conversations),
                 total_storage_mb=file_size_mb,
-                preferences=data.get('preferences', {}),
-                open_files=data.get('openFiles', [])[:100],
-                recent_files=data.get('recentFiles', [])[:100],
-                pinned_files=data.get('pinnedFiles', [])[:50]
+                preferences=data.get("preferences", {}),
+                open_files=data.get("openFiles", [])[:100],
+                recent_files=data.get("recentFiles", [])[:100],
+                pinned_files=data.get("pinnedFiles", [])[:50],
             )
-            
+
         except Exception as e:
             self.log_error(f"Failed to extract conversation state: {e}")
             return self._empty_conversation_state()
-    
-    def _create_conversation_context(self, 
-                                    conv: Dict[str, Any]) -> ConversationContext:
+
+    def _create_conversation_context(self, conv: Dict[str, Any]) -> ConversationContext:
         """Create ConversationContext from conversation data."""
         return ConversationContext(
-            conversation_id=conv.get('id', ''),
-            title=conv.get('title', 'Untitled'),
-            created_at=conv.get('createdAt', 0) / 1000,  # Convert from ms
-            updated_at=conv.get('updatedAt', 0) / 1000,
-            message_count=len(conv.get('messages', [])),
-            total_tokens=conv.get('totalTokens', 0),
-            max_tokens=conv.get('maxTokens', 100000),
+            conversation_id=conv.get("id", ""),
+            title=conv.get("title", "Untitled"),
+            created_at=conv.get("createdAt", 0) / 1000,  # Convert from ms
+            updated_at=conv.get("updatedAt", 0) / 1000,
+            message_count=len(conv.get("messages", [])),
+            total_tokens=conv.get("totalTokens", 0),
+            max_tokens=conv.get("maxTokens", 100000),
             referenced_files=[],  # Would need full extraction
-            open_tabs=conv.get('openTabs', []),
-            tags=conv.get('tags', []),
-            is_active=False
+            open_tabs=conv.get("openTabs", []),
+            tags=conv.get("tags", []),
+            is_active=False,
         )
-    
+
     def _empty_conversation_state(self) -> ConversationState:
         """Create empty conversation state."""
         return ConversationState(
@@ -462,76 +496,76 @@ class ContextPreservationService(BaseService):
             preferences={},
             open_files=[],
             recent_files=[],
-            pinned_files=[]
+            pinned_files=[],
         )
-    
+
     def _extract_file_paths(self, content: str) -> List[str]:
         """Extract file paths from message content."""
         import re
-        
+
         files = set()
-        
+
         # Common file path patterns
         patterns = [
             r'[\'"`]([/\\]?(?:[a-zA-Z]:)?[/\\]?[\w\-_./\\]+\.\w+)[\'"`]',
-            r'(?:^|\s)([/\\]?(?:[a-zA-Z]:)?[/\\]?[\w\-_./\\]+\.\w+)(?:\s|$)',
+            r"(?:^|\s)([/\\]?(?:[a-zA-Z]:)?[/\\]?[\w\-_./\\]+\.\w+)(?:\s|$)",
         ]
-        
+
         for pattern in patterns:
             matches = re.findall(pattern, content)
             files.update(matches)
-        
+
         return list(files)
-    
+
     async def _create_backup(self) -> Path:
         """Create backup of Claude configuration."""
         try:
             timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
             backup_name = f"claude_backup_{timestamp}.json"
-            
+
             # Compress if large
             file_size_mb = self.claude_json_path.stat().st_size / (1024 * 1024)
             if file_size_mb > self.compression_threshold_mb:
                 backup_name += ".gz"
                 backup_path = self.claude_backup_dir / backup_name
-                
-                with open(self.claude_json_path, 'rb') as f_in:
-                    with gzip.open(backup_path, 'wb') as f_out:
+
+                with open(self.claude_json_path, "rb") as f_in:
+                    with gzip.open(backup_path, "wb") as f_out:
                         shutil.copyfileobj(f_in, f_out)
             else:
                 backup_path = self.claude_backup_dir / backup_name
                 shutil.copy2(self.claude_json_path, backup_path)
-            
+
             self.log_info(f"Created backup: {backup_path}")
             return backup_path
-            
+
         except Exception as e:
             self.log_error(f"Failed to create backup: {e}")
             raise
-    
+
     def get_statistics(self) -> Dict[str, Any]:
         """Get context preservation statistics.
-        
+
         Returns:
             Dictionary containing statistics
         """
         claude_size_mb = 0.0
         if self.claude_json_path.exists():
             claude_size_mb = self.claude_json_path.stat().st_size / (1024 * 1024)
-        
+
         backup_count = 0
         backup_size_mb = 0.0
         if self.claude_backup_dir.exists():
             backups = list(self.claude_backup_dir.glob("claude_backup_*.json*"))
             backup_count = len(backups)
             backup_size_mb = sum(f.stat().st_size for f in backups) / (1024 * 1024)
-        
+
         return {
-            'claude_config_exists': self.claude_json_path.exists(),
-            'claude_config_size_mb': round(claude_size_mb, 2),
-            'is_large_file': claude_size_mb > self.large_file_threshold_mb,
-            'backup_count': backup_count,
-            'total_backup_size_mb': round(backup_size_mb, 2),
-            'files_processed': self.files_processed,
-            'total_size_processed_mb': round(self.total_size_processed_mb, 2)
-        }
+            "claude_config_exists": self.claude_json_path.exists(),
+            "claude_config_size_mb": round(claude_size_mb, 2),
+            "is_large_file": claude_size_mb > self.large_file_threshold_mb,
+            "backup_count": backup_count,
+            "total_backup_size_mb": round(backup_size_mb, 2),
+            "files_processed": self.files_processed,
+            "total_size_processed_mb": round(self.total_size_processed_mb, 2),
+        }