claude-mpm 3.5.4__py3-none-any.whl → 3.6.0__py3-none-any.whl

claude_mpm/models/agent_session.py

@@ -0,0 +1,511 @@
+"""Agent Session Model for Event Aggregation.
+
+WHY: This model represents a complete agent activity session, capturing all events
+from initial prompt through delegations, tool operations, and final responses.
+It provides a structured way to analyze what happened during an agent session.
+
+DESIGN DECISION: We use a hierarchical event structure to maintain relationships
+between related events (e.g., pre_tool and post_tool pairs) while preserving
+chronological order for session replay and analysis.
+"""
+
+import json
+import os
+from datetime import datetime
+from dataclasses import dataclass, field, asdict
+from typing import List, Dict, Any, Optional, Set
+from enum import Enum
+from pathlib import Path
+
+
+class EventCategory(Enum):
+    """Categories for different types of events in a session."""
+    PROMPT = "prompt"
+    DELEGATION = "delegation"
+    TOOL = "tool"
+    FILE = "file"
+    TODO = "todo"
+    RESPONSE = "response"
+    MEMORY = "memory"
+    STATUS = "status"
+    SYSTEM = "system"
+
+
+@dataclass
+class SessionEvent:
+    """Individual event within a session.
+    
+    WHY: Each event needs to be self-contained with all necessary context
+    for later analysis, including timing, category, and relationships.
+    """
+    timestamp: str
+    event_type: str  # Original event type from Socket.IO
+    category: EventCategory
+    data: Dict[str, Any]
+    session_id: Optional[str] = None
+    agent_context: Optional[str] = None  # Which agent was active
+    correlation_id: Optional[str] = None  # For matching pre/post events
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            'timestamp': self.timestamp,
+            'event_type': self.event_type,
+            'category': self.category.value,
+            'data': self.data,
+            'session_id': self.session_id,
+            'agent_context': self.agent_context,
+            'correlation_id': self.correlation_id
+        }
+
+
+@dataclass
+class ToolOperation:
+    """Represents a complete tool operation with pre/post events.
+    
+    WHY: Tool operations often span multiple events (pre_tool, post_tool).
+    This structure correlates them for complete analysis.
+    """
+    tool_name: str
+    agent_type: str
+    start_time: str
+    end_time: Optional[str] = None
+    input_data: Optional[Dict[str, Any]] = None
+    output_data: Optional[Dict[str, Any]] = None
+    duration_ms: Optional[int] = None
+    success: bool = True
+    error: Optional[str] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return asdict(self)
+
+
+@dataclass
+class AgentDelegation:
+    """Represents an agent delegation with its full lifecycle.
+    
+    WHY: Agent delegations are key session events that need special tracking
+    to understand the flow of work between agents.
+    """
+    agent_type: str
+    task_description: str
+    start_time: str
+    end_time: Optional[str] = None
+    prompt: Optional[str] = None
+    response: Optional[str] = None
+    tool_operations: List[ToolOperation] = field(default_factory=list)
+    file_changes: List[str] = field(default_factory=list)
+    todos_modified: List[Dict[str, Any]] = field(default_factory=list)
+    memory_updates: List[Dict[str, Any]] = field(default_factory=list)
+    duration_ms: Optional[int] = None
+    success: bool = True
+    error: Optional[str] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            'agent_type': self.agent_type,
+            'task_description': self.task_description,
+            'start_time': self.start_time,
+            'end_time': self.end_time,
+            'prompt': self.prompt,
+            'response': self.response,
+            'tool_operations': [op.to_dict() for op in self.tool_operations],
+            'file_changes': self.file_changes,
+            'todos_modified': self.todos_modified,
+            'memory_updates': self.memory_updates,
+            'duration_ms': self.duration_ms,
+            'success': self.success,
+            'error': self.error
+        }
+
+
+@dataclass
+class SessionMetrics:
+    """Aggregated metrics for a session.
+    
+    WHY: Quick summary statistics help identify patterns and anomalies
+    without processing all events.
+    """
+    total_events: int = 0
+    event_counts: Dict[str, int] = field(default_factory=dict)
+    agents_used: Set[str] = field(default_factory=set)
+    tools_used: Set[str] = field(default_factory=set)
+    files_modified: Set[str] = field(default_factory=set)
+    total_delegations: int = 0
+    total_tool_calls: int = 0
+    total_file_operations: int = 0
+    session_duration_ms: Optional[int] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            'total_events': self.total_events,
+            'event_counts': self.event_counts,
+            'agents_used': list(self.agents_used),
+            'tools_used': list(self.tools_used),
+            'files_modified': list(self.files_modified),
+            'total_delegations': self.total_delegations,
+            'total_tool_calls': self.total_tool_calls,
+            'total_file_operations': self.total_file_operations,
+            'session_duration_ms': self.session_duration_ms
+        }
+
+
+@dataclass
+class AgentSession:
+    """Complete representation of an agent activity session.
+    
+    WHY: This is the top-level model that captures everything that happened
+    during a Claude MPM session, from initial prompt to final response.
+    
+    DESIGN DECISION: We maintain both a flat chronological event list and
+    structured representations (delegations, tool operations) to support
+    different analysis needs.
+    """
+    session_id: str
+    start_time: str
+    end_time: Optional[str] = None
+    working_directory: str = ""
+    launch_method: str = ""
+    initial_prompt: Optional[str] = None
+    final_response: Optional[str] = None
+    
+    # Event collections
+    events: List[SessionEvent] = field(default_factory=list)
+    delegations: List[AgentDelegation] = field(default_factory=list)
+    
+    # Session state
+    current_agent: Optional[str] = None
+    active_delegation: Optional[AgentDelegation] = None
+    pending_tool_operations: Dict[str, ToolOperation] = field(default_factory=dict)
+    
+    # Metrics
+    metrics: SessionMetrics = field(default_factory=SessionMetrics)
+    
+    # Metadata
+    claude_pid: Optional[int] = None
+    git_branch: Optional[str] = None
+    project_root: Optional[str] = None
+    
+    def add_event(self, event_type: str, data: Dict[str, Any], timestamp: Optional[str] = None) -> SessionEvent:
+        """Add an event to the session.
+        
+        WHY: Centralizes event processing logic including categorization
+        and metric updates.
+        """
+        if timestamp is None:
+            timestamp = datetime.utcnow().isoformat() + 'Z'
+        
+        # Categorize the event
+        category = self._categorize_event(event_type, data)
+        
+        # Create the event
+        event = SessionEvent(
+            timestamp=timestamp,
+            event_type=event_type,
+            category=category,
+            data=data,
+            session_id=self.session_id,
+            agent_context=self.current_agent
+        )
+        
+        self.events.append(event)
+        
+        # Update metrics
+        self.metrics.total_events += 1
+        self.metrics.event_counts[event_type] = self.metrics.event_counts.get(event_type, 0) + 1
+        
+        # Process specific event types
+        self._process_event(event)
+        
+        return event
+    
+    def _categorize_event(self, event_type: str, data: Dict[str, Any]) -> EventCategory:
+        """Categorize an event based on its type and data.
+        
+        WHY: Categories help with filtering and analysis of related events.
+        """
+        # Check event type patterns
+        if 'prompt' in event_type.lower() or event_type == 'user_input':
+            return EventCategory.PROMPT
+        elif 'delegation' in event_type.lower() or event_type == 'Task':
+            return EventCategory.DELEGATION
+        elif 'tool' in event_type.lower() or event_type in ['PreToolUse', 'PostToolUse']:
+            return EventCategory.TOOL
+        elif 'file' in event_type.lower() or 'write' in event_type.lower() or 'read' in event_type.lower():
+            return EventCategory.FILE
+        elif 'todo' in event_type.lower():
+            return EventCategory.TODO
+        elif 'response' in event_type.lower() or event_type in ['Stop', 'SubagentStop']:
+            return EventCategory.RESPONSE
+        elif 'memory' in event_type.lower():
+            return EventCategory.MEMORY
+        elif 'status' in event_type.lower() or 'session' in event_type.lower():
+            return EventCategory.STATUS
+        else:
+            return EventCategory.SYSTEM
+    
+    def _process_event(self, event: SessionEvent):
+        """Process specific event types to update session state.
+        
+        WHY: Different event types require different processing to maintain
+        accurate session state and correlations.
+        """
+        event_type = event.event_type
+        data = event.data
+        
+        # Track user prompts
+        if event.category == EventCategory.PROMPT:
+            if not self.initial_prompt and 'prompt' in data:
+                self.initial_prompt = data['prompt']
+        
+        # Track agent delegations
+        elif event_type == 'Task' or 'delegation' in event_type.lower():
+            agent_type = data.get('agent_type', 'unknown')
+            self.current_agent = agent_type
+            self.metrics.agents_used.add(agent_type)
+            
+            # Create new delegation
+            delegation = AgentDelegation(
+                agent_type=agent_type,
+                task_description=data.get('description', ''),
+                start_time=event.timestamp,
+                prompt=data.get('prompt')
+            )
+            self.delegations.append(delegation)
+            self.active_delegation = delegation
+            self.metrics.total_delegations += 1
+        
+        # Track tool operations
+        elif event_type == 'PreToolUse':
+            tool_name = data.get('tool_name', 'unknown')
+            self.metrics.tools_used.add(tool_name)
+            self.metrics.total_tool_calls += 1
+            
+            # Create pending tool operation
+            tool_op = ToolOperation(
+                tool_name=tool_name,
+                agent_type=self.current_agent or 'unknown',
+                start_time=event.timestamp,
+                input_data=data.get('tool_input')
+            )
+            
+            # Store with correlation ID if available
+            correlation_id = f"{event.session_id}:{tool_name}:{event.timestamp}"
+            self.pending_tool_operations[correlation_id] = tool_op
+            event.correlation_id = correlation_id
+            
+            # Add to active delegation if exists
+            if self.active_delegation:
+                self.active_delegation.tool_operations.append(tool_op)
+        
+        elif event_type == 'PostToolUse':
+            # Match with pending tool operation
+            tool_name = data.get('tool_name', 'unknown')
+            
+            # Find matching pending operation
+            for corr_id, tool_op in list(self.pending_tool_operations.items()):
+                if tool_op.tool_name == tool_name and not tool_op.end_time:
+                    tool_op.end_time = event.timestamp
+                    tool_op.output_data = data.get('tool_output')
+                    tool_op.success = data.get('success', True)
+                    tool_op.error = data.get('error')
+                    
+                    # Calculate duration
+                    try:
+                        start = datetime.fromisoformat(tool_op.start_time.replace('Z', '+00:00'))
+                        end = datetime.fromisoformat(event.timestamp.replace('Z', '+00:00'))
+                        tool_op.duration_ms = int((end - start).total_seconds() * 1000)
+                    except:
+                        pass
+                    
+                    event.correlation_id = corr_id
+                    del self.pending_tool_operations[corr_id]
+                    break
+        
+        # Track file operations
+        elif event.category == EventCategory.FILE:
+            file_path = data.get('file_path') or data.get('path') or data.get('file')
+            if file_path:
+                self.metrics.files_modified.add(file_path)
+                self.metrics.total_file_operations += 1
+                
+                if self.active_delegation:
+                    self.active_delegation.file_changes.append(file_path)
+        
+        # Track responses
+        elif event_type in ['Stop', 'SubagentStop']:
+            response = data.get('response') or data.get('content') or data.get('message')
+            if response:
+                if event_type == 'SubagentStop' and self.active_delegation:
+                    self.active_delegation.response = response
+                    self.active_delegation.end_time = event.timestamp
+                    self.active_delegation = None
+                elif event_type == 'Stop':
+                    self.final_response = response
+        
+        # Track todo updates
+        elif event.category == EventCategory.TODO:
+            if self.active_delegation and 'todos' in data:
+                self.active_delegation.todos_modified.append(data['todos'])
+        
+        # Track memory updates
+        elif event.category == EventCategory.MEMORY:
+            if self.active_delegation:
+                self.active_delegation.memory_updates.append(data)
+    
+    def finalize(self):
+        """Finalize the session by calculating final metrics.
+        
+        WHY: Some metrics can only be calculated after all events are processed.
+        """
+        if not self.end_time and self.events:
+            self.end_time = self.events[-1].timestamp
+        
+        # Calculate session duration
+        if self.start_time and self.end_time:
+            try:
+                start = datetime.fromisoformat(self.start_time.replace('Z', '+00:00'))
+                end = datetime.fromisoformat(self.end_time.replace('Z', '+00:00'))
+                self.metrics.session_duration_ms = int((end - start).total_seconds() * 1000)
+            except:
+                pass
+        
+        # Finalize any pending delegations
+        for delegation in self.delegations:
+            if not delegation.end_time:
+                delegation.end_time = self.end_time
+                delegation.success = False
+                delegation.error = "Delegation did not complete"
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            'session_id': self.session_id,
+            'start_time': self.start_time,
+            'end_time': self.end_time,
+            'working_directory': self.working_directory,
+            'launch_method': self.launch_method,
+            'initial_prompt': self.initial_prompt,
+            'final_response': self.final_response,
+            'events': [e.to_dict() for e in self.events],
+            'delegations': [d.to_dict() for d in self.delegations],
+            'metrics': self.metrics.to_dict(),
+            'metadata': {
+                'claude_pid': self.claude_pid,
+                'git_branch': self.git_branch,
+                'project_root': self.project_root
+            }
+        }
+    
+    def save_to_file(self, directory: Optional[str] = None) -> str:
+        """Save the session to a JSON file.
+        
+        WHY: Persistent storage allows for later analysis and debugging.
+        
+        Args:
+            directory: Directory to save to (defaults to .claude-mpm/sessions/)
+            
+        Returns:
+            Path to the saved file
+        """
+        if directory is None:
+            directory = Path.cwd() / '.claude-mpm' / 'sessions'
+        else:
+            directory = Path(directory)
+        
+        # Create directory if it doesn't exist
+        directory.mkdir(parents=True, exist_ok=True)
+        
+        # Generate filename with timestamp
+        timestamp = self.start_time.replace(':', '-').replace('.', '-')[:19]
+        filename = f"session_{self.session_id[:8]}_{timestamp}.json"
+        filepath = directory / filename
+        
+        # Save to file
+        with open(filepath, 'w', encoding='utf-8') as f:
+            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)
+        
+        return str(filepath)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'AgentSession':
+        """Create an AgentSession from a dictionary.
+        
+        WHY: Allows loading saved sessions for analysis.
+        """
+        session = cls(
+            session_id=data['session_id'],
+            start_time=data['start_time'],
+            end_time=data.get('end_time'),
+            working_directory=data.get('working_directory', ''),
+            launch_method=data.get('launch_method', ''),
+            initial_prompt=data.get('initial_prompt'),
+            final_response=data.get('final_response')
+        )
+        
+        # Restore events
+        for event_data in data.get('events', []):
+            event = SessionEvent(
+                timestamp=event_data['timestamp'],
+                event_type=event_data['event_type'],
+                category=EventCategory(event_data['category']),
+                data=event_data['data'],
+                session_id=event_data.get('session_id'),
+                agent_context=event_data.get('agent_context'),
+                correlation_id=event_data.get('correlation_id')
+            )
+            session.events.append(event)
+        
+        # Restore delegations
+        for del_data in data.get('delegations', []):
+            delegation = AgentDelegation(
+                agent_type=del_data['agent_type'],
+                task_description=del_data['task_description'],
+                start_time=del_data['start_time'],
+                end_time=del_data.get('end_time'),
+                prompt=del_data.get('prompt'),
+                response=del_data.get('response'),
+                tool_operations=[ToolOperation(**op) for op in del_data.get('tool_operations', [])],
+                file_changes=del_data.get('file_changes', []),
+                todos_modified=del_data.get('todos_modified', []),
+                memory_updates=del_data.get('memory_updates', []),
+                duration_ms=del_data.get('duration_ms'),
+                success=del_data.get('success', True),
+                error=del_data.get('error')
+            )
+            session.delegations.append(delegation)
+        
+        # Restore metrics
+        metrics_data = data.get('metrics', {})
+        session.metrics = SessionMetrics(
+            total_events=metrics_data.get('total_events', 0),
+            event_counts=metrics_data.get('event_counts', {}),
+            agents_used=set(metrics_data.get('agents_used', [])),
+            tools_used=set(metrics_data.get('tools_used', [])),
+            files_modified=set(metrics_data.get('files_modified', [])),
+            total_delegations=metrics_data.get('total_delegations', 0),
+            total_tool_calls=metrics_data.get('total_tool_calls', 0),
+            total_file_operations=metrics_data.get('total_file_operations', 0),
+            session_duration_ms=metrics_data.get('session_duration_ms')
+        )
+        
+        # Restore metadata
+        metadata = data.get('metadata', {})
+        session.claude_pid = metadata.get('claude_pid')
+        session.git_branch = metadata.get('git_branch')
+        session.project_root = metadata.get('project_root')
+        
+        return session
+    
+    @classmethod
+    def load_from_file(cls, filepath: str) -> 'AgentSession':
+        """Load a session from a JSON file.
+        
+        WHY: Enables analysis of historical sessions.
+        """
+        with open(filepath, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        return cls.from_dict(data)