MockLoop MCP Server

""" MCP Context Management System for MockLoop This module provides comprehensive context management for stateful testing workflows, cross-session orchestration, and agent-specific state management. Features: - Multiple context types for different workflow scenarios - Context persistence and storage using SQLite - Context inheritance and hierarchical relationships - Context snapshots for rollback capabilities - Cross-session data sharing via GlobalContext - Agent-specific state management via AgentContext - Thread-safe operations for multi-agent scenarios - Comprehensive audit logging integration """ import asyncio import json import sqlite3 import threading import time import uuid from contextlib import asynccontextmanager from datetime import datetime, timedelta from enum import Enum from pathlib import Path from typing import Any, Optional, Union import logging # Configure logger for this module logger = logging.getLogger(__name__) class ContextType(Enum): """Enumeration of available context types.""" TEST_SESSION = "test_session" WORKFLOW = "workflow" SCENARIO = "scenario" PERFORMANCE = "performance" GLOBAL = "global" AGENT = "agent" class ContextStatus(Enum): """Enumeration of context status values.""" ACTIVE = "active" PAUSED = "paused" COMPLETED = "completed" EXPIRED = "expired" ERROR = "error" class ContextError(Exception): """Base exception for context management errors.""" pass class ContextNotFoundError(ContextError): """Raised when a requested context is not found.""" pass class ContextValidationError(ContextError): """Raised when context validation fails.""" pass class ContextPersistenceError(ContextError): """Raised when context persistence operations fail.""" pass class ContextSnapshot: """Represents a snapshot of context state for rollback capabilities.""" def __init__( self, context_id: str, snapshot_data: dict[str, Any], timestamp: datetime | None = None, ): self.context_id = context_id self.snapshot_data = snapshot_data self.timestamp = timestamp or datetime.utcnow() self.snapshot_id = str(uuid.uuid4()) def to_dict(self) -> dict[str, Any]: """Convert snapshot to dictionary representation.""" return { "snapshot_id": self.snapshot_id, "context_id": self.context_id, "snapshot_data": self.snapshot_data, "timestamp": self.timestamp.isoformat(), } @classmethod def from_dict(cls, data: dict[str, Any]) -> "ContextSnapshot": """Create snapshot from dictionary representation.""" return cls( context_id=data["context_id"], snapshot_data=data["snapshot_data"], timestamp=datetime.fromisoformat(data["timestamp"]), ) class BaseContext: """Base class for all context types.""" def __init__( self, context_id: str, context_type: ContextType, data: dict[str, Any] | None = None, parent_context_id: str | None = None, metadata: dict[str, Any] | None = None, ): self.context_id = context_id self.context_type = context_type self.data = data or {} self.parent_context_id = parent_context_id self.metadata = metadata or {} self.created_at = datetime.utcnow() self.updated_at = self.created_at self.status = ContextStatus.ACTIVE self.expiry_time: datetime | None = None self.tags: set[str] = set() self._lock = threading.RLock() self._snapshots: list[ContextSnapshot] = [] def update_data(self, updates: dict[str, Any], merge: bool = True) -> None: """Update context data with change tracking.""" with self._lock: if merge: self.data.update(updates) else: self.data = updates.copy() self.updated_at = datetime.utcnow() def get_data(self, key: str | None = None, default: Any = None) -> Any: """Get context data with optional key filtering.""" with self._lock: if key is None: return self.data.copy() return self.data.get(key, default) def add_tag(self, tag: str) -> None: """Add a tag to the context.""" with self._lock: self.tags.add(tag) self.updated_at = datetime.utcnow() def remove_tag(self, tag: str) -> None: """Remove a tag from the context.""" with self._lock: self.tags.discard(tag) self.updated_at = datetime.utcnow() def has_tag(self, tag: str) -> bool: """Check if context has a specific tag.""" with self._lock: return tag in self.tags def set_expiry(self, expiry_time: datetime) -> None: """Set context expiry time.""" with self._lock: self.expiry_time = expiry_time self.updated_at = datetime.utcnow() def is_expired(self) -> bool: """Check if context has expired.""" with self._lock: if self.expiry_time is None: return False return datetime.utcnow() > self.expiry_time def create_snapshot(self, description: str | None = None) -> ContextSnapshot: """Create a snapshot of current context state.""" with self._lock: snapshot_data = { "data": self.data.copy(), "metadata": self.metadata.copy(), "status": self.status.value, "tags": list(self.tags), "description": description, } snapshot = ContextSnapshot(self.context_id, snapshot_data) self._snapshots.append(snapshot) return snapshot def restore_from_snapshot(self, snapshot_id: str) -> bool: """Restore context state from a snapshot.""" with self._lock: for snapshot in self._snapshots: if snapshot.snapshot_id == snapshot_id: snapshot_data = snapshot.snapshot_data self.data = snapshot_data.get("data", {}).copy() self.metadata = snapshot_data.get("metadata", {}).copy() self.status = ContextStatus(snapshot_data.get("status", "active")) self.tags = set(snapshot_data.get("tags", [])) self.updated_at = datetime.utcnow() return True return False def list_snapshots(self) -> list[dict[str, Any]]: """List all snapshots for this context.""" with self._lock: return [snapshot.to_dict() for snapshot in self._snapshots] def to_dict(self) -> dict[str, Any]: """Convert context to dictionary representation.""" with self._lock: return { "context_id": self.context_id, "context_type": self.context_type.value, "data": self.data.copy(), "parent_context_id": self.parent_context_id, "metadata": self.metadata.copy(), "created_at": self.created_at.isoformat(), "updated_at": self.updated_at.isoformat(), "status": self.status.value, "expiry_time": self.expiry_time.isoformat() if self.expiry_time else None, "tags": list(self.tags), "snapshots": [snapshot.to_dict() for snapshot in self._snapshots], } @classmethod def from_dict(cls, data: dict[str, Any]) -> "BaseContext": """Create context from dictionary representation.""" context = cls( context_id=data["context_id"], context_type=ContextType(data["context_type"]), data=data.get("data", {}), parent_context_id=data.get("parent_context_id"), metadata=data.get("metadata", {}), ) context.created_at = datetime.fromisoformat(data["created_at"]) context.updated_at = datetime.fromisoformat(data["updated_at"]) context.status = ContextStatus(data["status"]) if data.get("expiry_time"): context.expiry_time = datetime.fromisoformat(data["expiry_time"]) context.tags = set(data.get("tags", [])) # Restore snapshots for snapshot_data in data.get("snapshots", []): snapshot = ContextSnapshot.from_dict(snapshot_data) context._snapshots.append(snapshot) return context class TestSessionContext(BaseContext): """Context for managing test session state and metadata.""" def __init__( self, context_id: str, session_name: str, test_plan: dict[str, Any], session_config: dict[str, Any] | None = None, ): super().__init__(context_id, ContextType.TEST_SESSION) self.session_name = session_name self.test_plan = test_plan self.session_config = session_config or {} self.test_results: list[dict[str, Any]] = [] self.current_test_index = 0 # Initialize session-specific data self.data.update( { "session_name": session_name, "test_plan": test_plan, "session_config": self.session_config, "test_results": self.test_results, "current_test_index": self.current_test_index, "session_metrics": { "total_tests": len(test_plan.get("tests", [])), "completed_tests": 0, "failed_tests": 0, "start_time": self.created_at.isoformat(), }, } ) def add_test_result(self, test_result: dict[str, Any]) -> None: """Add a test result to the session.""" with self._lock: self.test_results.append(test_result) self.data["test_results"] = self.test_results # Update metrics metrics = self.data["session_metrics"] metrics["completed_tests"] = len(self.test_results) if test_result.get("status") == "failed": metrics["failed_tests"] += 1 self.updated_at = datetime.utcnow() def get_session_summary(self) -> dict[str, Any]: """Get a summary of the test session.""" with self._lock: metrics = self.data["session_metrics"] duration = (self.updated_at - self.created_at).total_seconds() return { "session_id": self.context_id, "session_name": self.session_name, "status": self.status.value, "duration_seconds": duration, "total_tests": metrics["total_tests"], "completed_tests": metrics["completed_tests"], "failed_tests": metrics["failed_tests"], "success_rate": (metrics["completed_tests"] - metrics["failed_tests"]) / max(metrics["completed_tests"], 1), "created_at": self.created_at.isoformat(), "updated_at": self.updated_at.isoformat(), } class WorkflowContext(BaseContext): """Context for tracking multi-step testing workflows.""" def __init__( self, context_id: str, workflow_name: str, workflow_steps: list[dict[str, Any]] ): super().__init__(context_id, ContextType.WORKFLOW) self.workflow_name = workflow_name self.workflow_steps = workflow_steps self.current_step_index = 0 self.step_results: list[dict[str, Any]] = [] # Initialize workflow-specific data self.data.update( { "workflow_name": workflow_name, "workflow_steps": workflow_steps, "current_step_index": self.current_step_index, "step_results": self.step_results, "workflow_state": "initialized", "can_pause": True, "can_resume": True, } ) def advance_step(self, step_result: dict[str, Any] | None = None) -> bool: """Advance to the next workflow step.""" with self._lock: if step_result: self.step_results.append(step_result) self.data["step_results"] = self.step_results if self.current_step_index < len(self.workflow_steps) - 1: self.current_step_index += 1 self.data["current_step_index"] = self.current_step_index self.updated_at = datetime.utcnow() return True else: self.data["workflow_state"] = "completed" self.status = ContextStatus.COMPLETED return False def get_current_step(self) -> dict[str, Any] | None: """Get the current workflow step.""" with self._lock: if self.current_step_index < len(self.workflow_steps): return self.workflow_steps[self.current_step_index] return None def pause_workflow(self) -> bool: """Pause the workflow execution.""" with self._lock: if self.data.get("can_pause", True): self.status = ContextStatus.PAUSED self.data["workflow_state"] = "paused" self.updated_at = datetime.utcnow() return True return False def resume_workflow(self) -> bool: """Resume the workflow execution.""" with self._lock: if self.status == ContextStatus.PAUSED and self.data.get( "can_resume", True ): self.status = ContextStatus.ACTIVE self.data["workflow_state"] = "running" self.updated_at = datetime.utcnow() return True return False class ScenarioContext(BaseContext): """Context for managing active scenario configurations.""" def __init__( self, context_id: str, scenario_name: str, scenario_config: dict[str, Any] ): super().__init__(context_id, ContextType.SCENARIO) self.scenario_name = scenario_name self.scenario_config = scenario_config self.deployment_info: dict[str, Any] | None = None # Initialize scenario-specific data self.data.update( { "scenario_name": scenario_name, "scenario_config": scenario_config, "deployment_info": self.deployment_info, "scenario_state": "configured", "active_endpoints": [], "scenario_metrics": { "requests_processed": 0, "errors_generated": 0, "last_activity": None, }, } ) def deploy_scenario(self, deployment_info: dict[str, Any]) -> None: """Mark scenario as deployed with deployment information.""" with self._lock: self.deployment_info = deployment_info self.data["deployment_info"] = deployment_info self.data["scenario_state"] = "deployed" self.updated_at = datetime.utcnow() def update_metrics(self, metric_updates: dict[str, Any]) -> None: """Update scenario metrics.""" with self._lock: metrics = self.data["scenario_metrics"] metrics.update(metric_updates) metrics["last_activity"] = datetime.utcnow().isoformat() self.updated_at = datetime.utcnow() class PerformanceContext(BaseContext): """Context for tracking performance metrics across test runs.""" def __init__(self, context_id: str, test_run_id: str): super().__init__(context_id, ContextType.PERFORMANCE) self.test_run_id = test_run_id self.metrics_history: list[dict[str, Any]] = [] # Initialize performance-specific data self.data.update( { "test_run_id": test_run_id, "metrics_history": self.metrics_history, "current_metrics": {}, "performance_thresholds": { "max_response_time_ms": 1000, "min_throughput_rps": 100, "max_error_rate_percent": 5.0, }, "alerts": [], } ) def add_metrics(self, metrics: dict[str, Any]) -> None: """Add performance metrics to the context.""" with self._lock: timestamped_metrics = { **metrics, "timestamp": datetime.utcnow().isoformat(), } self.metrics_history.append(timestamped_metrics) self.data["metrics_history"] = self.metrics_history self.data["current_metrics"] = metrics # Check thresholds and generate alerts self._check_performance_thresholds(metrics) self.updated_at = datetime.utcnow() def _check_performance_thresholds(self, metrics: dict[str, Any]) -> None: """Check performance metrics against thresholds and generate alerts.""" thresholds = self.data["performance_thresholds"] alerts = self.data["alerts"] # Check response time if metrics.get("avg_response_time_ms", 0) > thresholds["max_response_time_ms"]: alerts.append( { "type": "response_time_exceeded", "message": f"Average response time {metrics['avg_response_time_ms']}ms exceeds threshold {thresholds['max_response_time_ms']}ms", "timestamp": datetime.utcnow().isoformat(), "severity": "warning", } ) # Check throughput if metrics.get("throughput_rps", 0) < thresholds["min_throughput_rps"]: alerts.append( { "type": "throughput_below_threshold", "message": f"Throughput {metrics['throughput_rps']} RPS below threshold {thresholds['min_throughput_rps']} RPS", "timestamp": datetime.utcnow().isoformat(), "severity": "warning", } ) # Check error rate if metrics.get("error_rate_percent", 0) > thresholds["max_error_rate_percent"]: alerts.append( { "type": "error_rate_exceeded", "message": f"Error rate {metrics['error_rate_percent']}% exceeds threshold {thresholds['max_error_rate_percent']}%", "timestamp": datetime.utcnow().isoformat(), "severity": "critical", } ) class GlobalContext(BaseContext): """Context for data shared across sessions and AI-driven orchestration.""" def __init__(self, context_id: str = "global_context"): super().__init__(context_id, ContextType.GLOBAL) # Initialize global-specific data self.data.update( { "shared_configurations": {}, "cross_session_metrics": { "total_sessions": 0, "total_tests_executed": 0, "global_success_rate": 0.0, "most_common_failures": [], "performance_trends": [], }, "global_settings": { "default_timeout_seconds": 300, "max_concurrent_sessions": 10, "enable_cross_session_learning": True, "auto_optimize_scenarios": True, }, "orchestration_state": { "active_orchestrations": [], "queued_orchestrations": [], "orchestration_history": [], }, "shared_test_data": {}, "global_analytics": { "api_usage_patterns": {}, "performance_baselines": {}, "security_findings": [], }, } ) def update_cross_session_metrics(self, session_metrics: dict[str, Any]) -> None: """Update global metrics with data from a completed session.""" with self._lock: global_metrics = self.data["cross_session_metrics"] global_metrics["total_sessions"] += 1 global_metrics["total_tests_executed"] += session_metrics.get( "total_tests", 0 ) # Update global success rate total_successful = global_metrics[ "total_tests_executed" ] - session_metrics.get("failed_tests", 0) global_metrics["global_success_rate"] = total_successful / max( global_metrics["total_tests_executed"], 1 ) self.updated_at = datetime.utcnow() def add_orchestration(self, orchestration_config: dict[str, Any]) -> str: """Add a new orchestration to the queue.""" with self._lock: orchestration_id = str(uuid.uuid4()) orchestration = { "orchestration_id": orchestration_id, "config": orchestration_config, "status": "queued", "created_at": datetime.utcnow().isoformat(), } self.data["orchestration_state"]["queued_orchestrations"].append( orchestration ) self.updated_at = datetime.utcnow() return orchestration_id def start_orchestration(self, orchestration_id: str) -> bool: """Move an orchestration from queued to active.""" with self._lock: queued = self.data["orchestration_state"]["queued_orchestrations"] active = self.data["orchestration_state"]["active_orchestrations"] for i, orchestration in enumerate(queued): if orchestration["orchestration_id"] == orchestration_id: orchestration["status"] = "active" orchestration["started_at"] = datetime.utcnow().isoformat() active.append(orchestration) queued.pop(i) self.updated_at = datetime.utcnow() return True return False class AgentContext(BaseContext): """Context for integrating agent-specific state when using LangGraph/CrewAI.""" def __init__( self, context_id: str, agent_id: str, agent_type: str, agent_config: dict[str, Any] | None = None, ): super().__init__(context_id, ContextType.AGENT) self.agent_id = agent_id self.agent_type = agent_type # "langgraph", "crewai", "custom" self.agent_config = agent_config or {} # Initialize agent-specific data self.data.update( { "agent_id": agent_id, "agent_type": agent_type, "agent_config": self.agent_config, "agent_state": {}, "memory": {"short_term": {}, "long_term": {}, "learned_patterns": []}, "workflow_state": { "current_workflow": None, "workflow_history": [], "pending_actions": [], }, "collaboration": { "connected_agents": [], "shared_context_ids": [], "communication_log": [], }, "performance_tracking": { "tasks_completed": 0, "success_rate": 0.0, "average_task_duration": 0.0, "optimization_suggestions": [], }, } ) def update_agent_state(self, state_updates: dict[str, Any]) -> None: """Update agent-specific state.""" with self._lock: self.data["agent_state"].update(state_updates) self.updated_at = datetime.utcnow() def add_memory(self, memory_type: str, key: str, value: Any) -> None: """Add information to agent memory.""" with self._lock: if memory_type in self.data["memory"]: self.data["memory"][memory_type][key] = value self.updated_at = datetime.utcnow() def get_memory(self, memory_type: str, key: str | None = None) -> Any: """Retrieve information from agent memory.""" with self._lock: memory = self.data["memory"].get(memory_type, {}) if key is None: return memory.copy() return memory.get(key) def add_learned_pattern(self, pattern: dict[str, Any]) -> None: """Add a learned pattern to agent memory.""" with self._lock: self.data["memory"]["learned_patterns"].append( {**pattern, "learned_at": datetime.utcnow().isoformat()} ) self.updated_at = datetime.utcnow() def connect_agent(self, other_agent_id: str, shared_context_id: str) -> None: """Connect this agent to another agent for collaboration.""" with self._lock: collaboration = self.data["collaboration"] if other_agent_id not in collaboration["connected_agents"]: collaboration["connected_agents"].append(other_agent_id) if shared_context_id not in collaboration["shared_context_ids"]: collaboration["shared_context_ids"].append(shared_context_id) self.updated_at = datetime.utcnow() def log_communication(self, message: dict[str, Any]) -> None: """Log communication with other agents.""" with self._lock: self.data["collaboration"]["communication_log"].append( {**message, "timestamp": datetime.utcnow().isoformat()} ) self.updated_at = datetime.utcnow() def update_performance_metrics(self, task_result: dict[str, Any]) -> None: """Update agent performance tracking.""" with self._lock: performance = self.data["performance_tracking"] performance["tasks_completed"] += 1 if task_result.get("success", False): # Update success rate total_tasks = performance["tasks_completed"] current_successes = performance["success_rate"] * (total_tasks - 1) performance["success_rate"] = (current_successes + 1) / total_tasks # Update average task duration if "duration_seconds" in task_result: total_tasks = performance["tasks_completed"] current_avg = performance["average_task_duration"] new_duration = task_result["duration_seconds"] performance["average_task_duration"] = ( current_avg * (total_tasks - 1) + new_duration ) / total_tasks self.updated_at = datetime.utcnow() class ContextStorage: """Handles context persistence and storage using SQLite database.""" def __init__(self, db_path: str = "mcp_context.db"): self.db_path = db_path self._lock = threading.RLock() self._init_database() def _init_database(self) -> None: """Initialize the SQLite database with required tables.""" with self._lock: conn = sqlite3.connect(self.db_path) try: cursor = conn.cursor() # Create contexts table cursor.execute(""" CREATE TABLE IF NOT EXISTS contexts ( context_id TEXT PRIMARY KEY, context_type TEXT NOT NULL, parent_context_id TEXT, data TEXT NOT NULL, metadata TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, status TEXT NOT NULL, expiry_time TEXT, tags TEXT, FOREIGN KEY (parent_context_id) REFERENCES contexts (context_id) ) """) # Create snapshots table cursor.execute(""" CREATE TABLE IF NOT EXISTS context_snapshots ( snapshot_id TEXT PRIMARY KEY, context_id TEXT NOT NULL, snapshot_data TEXT NOT NULL, timestamp TEXT NOT NULL, FOREIGN KEY (context_id) REFERENCES contexts (context_id) ) """) # Create indexes for better performance cursor.execute( "CREATE INDEX IF NOT EXISTS idx_context_type ON contexts (context_type)" ) cursor.execute( "CREATE INDEX IF NOT EXISTS idx_context_status ON contexts (status)" ) cursor.execute( "CREATE INDEX IF NOT EXISTS idx_context_parent ON contexts (parent_context_id)" ) cursor.execute( "CREATE INDEX IF NOT EXISTS idx_snapshot_context ON context_snapshots (context_id)" ) conn.commit() finally: conn.close() def save_context(self, context: BaseContext) -> None: """Save context to database.""" with self._lock: conn = sqlite3.connect(self.db_path) try: cursor = conn.cursor() cursor.execute( """ INSERT OR REPLACE INTO contexts (context_id, context_type, parent_context_id, data, metadata, created_at, updated_at, status, expiry_time, tags) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( context.context_id, context.context_type.value, context.parent_context_id, json.dumps(context.data), json.dumps(context.metadata), context.created_at.isoformat(), context.updated_at.isoformat(), context.status.value, context.expiry_time.isoformat() if context.expiry_time else None, json.dumps(list(context.tags)), ), ) # Save snapshots for snapshot in context._snapshots: cursor.execute( """ INSERT OR REPLACE INTO context_snapshots (snapshot_id, context_id, snapshot_data, timestamp) VALUES (?, ?, ?, ?) """, ( snapshot.snapshot_id, snapshot.context_id, json.dumps(snapshot.snapshot_data), snapshot.timestamp.isoformat(), ), ) conn.commit() finally: conn.close() def load_context(self, context_id: str) -> BaseContext | None: """Load context from database.""" with self._lock: conn = sqlite3.connect(self.db_path) try: cursor = conn.cursor() cursor.execute( """ SELECT context_id, context_type, parent_context_id, data, metadata, created_at, updated_at, status, expiry_time, tags FROM contexts WHERE context_id = ? """, (context_id,), ) row = cursor.fetchone() if not row: return None # Load snapshots cursor.execute( """ SELECT snapshot_id, snapshot_data, timestamp FROM context_snapshots WHERE context_id = ? ORDER BY timestamp """, (context_id,), ) snapshots = [] for snapshot_row in cursor.fetchall(): snapshot_data = { "snapshot_id": snapshot_row[0], "context_id": context_id, "snapshot_data": json.loads(snapshot_row[1]), "timestamp": snapshot_row[2], } snapshots.append(ContextSnapshot.from_dict(snapshot_data)) # Create context object based on type context_data = { "context_id": row[0], "context_type": row[1], "parent_context_id": row[2], "data": json.loads(row[3]), "metadata": json.loads(row[4]) if row[4] else {}, "created_at": row[5], "updated_at": row[6], "status": row[7], "expiry_time": row[8], "tags": json.loads(row[9]) if row[9] else [], "snapshots": [snapshot.to_dict() for snapshot in snapshots], } # Create appropriate context type context_type = ContextType(row[1]) if context_type == ContextType.TEST_SESSION: context = TestSessionContext( context_id=row[0], session_name=context_data["data"].get("session_name", ""), test_plan=context_data["data"].get("test_plan", {}), session_config=context_data["data"].get("session_config", {}), ) elif context_type == ContextType.WORKFLOW: context = WorkflowContext( context_id=row[0], workflow_name=context_data["data"].get("workflow_name", ""), workflow_steps=context_data["data"].get("workflow_steps", []), ) elif context_type == ContextType.SCENARIO: context = ScenarioContext( context_id=row[0], scenario_name=context_data["data"].get("scenario_name", ""), scenario_config=context_data["data"].get("scenario_config", {}), ) elif context_type == ContextType.PERFORMANCE: context = PerformanceContext( context_id=row[0], test_run_id=context_data["data"].get("test_run_id", ""), ) elif context_type == ContextType.GLOBAL: context = GlobalContext(context_id=row[0]) elif context_type == ContextType.AGENT: context = AgentContext( context_id=row[0], agent_id=context_data["data"].get("agent_id", ""), agent_type=context_data["data"].get("agent_type", ""), agent_config=context_data["data"].get("agent_config", {}), ) else: context = BaseContext.from_dict(context_data) # Restore state from loaded data context.data = context_data["data"] context.metadata = context_data["metadata"] context.parent_context_id = context_data["parent_context_id"] context.created_at = datetime.fromisoformat(context_data["created_at"]) context.updated_at = datetime.fromisoformat(context_data["updated_at"]) context.status = ContextStatus(context_data["status"]) if context_data["expiry_time"]: context.expiry_time = datetime.fromisoformat( context_data["expiry_time"] ) context.tags = set(context_data["tags"]) context._snapshots = snapshots return context finally: conn.close() def list_contexts( self, context_type: ContextType | None = None, status: ContextStatus | None = None, parent_context_id: str | None = None, tags: list[str] | None = None, ) -> list[dict[str, Any]]: """List contexts with optional filtering.""" with self._lock: conn = sqlite3.connect(self.db_path) try: cursor = conn.cursor() query = """ SELECT context_id, context_type, parent_context_id, data, metadata, created_at, updated_at, status, expiry_time, tags FROM contexts WHERE 1=1 """ params = [] if context_type: query += " AND context_type = ?" params.append(context_type.value) if status: query += " AND status = ?" params.append(status.value) if parent_context_id: query += " AND parent_context_id = ?" params.append(parent_context_id) query += " ORDER BY created_at DESC" cursor.execute(query, params) contexts = [] for row in cursor.fetchall(): context_tags = set(json.loads(row[9]) if row[9] else []) # Filter by tags if specified if tags and not any(tag in context_tags for tag in tags): continue contexts.append( { "context_id": row[0], "context_type": row[1], "parent_context_id": row[2], "data": json.loads(row[3]), "metadata": json.loads(row[4]) if row[4] else {}, "created_at": row[5], "updated_at": row[6], "status": row[7], "expiry_time": row[8], "tags": list(context_tags), } ) return contexts finally: conn.close() def delete_context(self, context_id: str) -> bool: """Delete context and its snapshots from database.""" with self._lock: conn = sqlite3.connect(self.db_path) try: cursor = conn.cursor() # Delete snapshots first (foreign key constraint) cursor.execute( "DELETE FROM context_snapshots WHERE context_id = ?", (context_id,) ) # Delete context cursor.execute( "DELETE FROM contexts WHERE context_id = ?", (context_id,) ) deleted = cursor.rowcount > 0 conn.commit() return deleted finally: conn.close() def cleanup_expired_contexts(self) -> int: """Remove expired contexts from database.""" with self._lock: conn = sqlite3.connect(self.db_path) try: cursor = conn.cursor() current_time = datetime.utcnow().isoformat() # Find expired contexts cursor.execute( """ SELECT context_id FROM contexts WHERE expiry_time IS NOT NULL AND expiry_time < ? """, (current_time,), ) expired_ids = [row[0] for row in cursor.fetchall()] # Delete expired contexts and their snapshots for context_id in expired_ids: cursor.execute( "DELETE FROM context_snapshots WHERE context_id = ?", (context_id,), ) cursor.execute( "DELETE FROM contexts WHERE context_id = ?", (context_id,) ) conn.commit() return len(expired_ids) finally: conn.close() class ContextManager: """Main context manager for coordinating all context operations.""" def __init__( self, db_path: str = "mcp_context.db", audit_logger: Any | None = None ): self.storage = ContextStorage(db_path) self.audit_logger = audit_logger self._contexts: dict[str, BaseContext] = {} self._lock = threading.RLock() self._global_context: GlobalContext | None = None # Initialize global context self._init_global_context() # Start cleanup task self._cleanup_task = None self._start_cleanup_task() def _init_global_context(self) -> None: """Initialize or load the global context.""" try: global_context = self.storage.load_context("global_context") if global_context and isinstance(global_context, GlobalContext): self._global_context = global_context self._contexts["global_context"] = global_context else: # Create new global context self._global_context = GlobalContext() self.storage.save_context(self._global_context) self._contexts["global_context"] = self._global_context except Exception: logger.exception("Failed to initialize global context") # Create new global context as fallback self._global_context = GlobalContext() self._contexts["global_context"] = self._global_context def _start_cleanup_task(self) -> None: """Start background task for cleaning up expired contexts.""" def cleanup_worker(): while True: try: time.sleep(300) # Run every 5 minutes self.cleanup_expired_contexts() except Exception: logger.exception("Error in cleanup task") self._cleanup_task = threading.Thread(target=cleanup_worker, daemon=True) self._cleanup_task.start() def create_context( self, context_type: ContextType, context_id: str | None = None, parent_context_id: str | None = None, **kwargs, ) -> BaseContext: """Create a new context of the specified type.""" with self._lock: if context_id is None: context_id = str(uuid.uuid4()) # Create context based on type if context_type == ContextType.TEST_SESSION: context = TestSessionContext( context_id=context_id, session_name=kwargs.get( "session_name", f"session_{context_id[:8]}" ), test_plan=kwargs.get("test_plan", {}), session_config=kwargs.get("session_config", {}), ) elif context_type == ContextType.WORKFLOW: context = WorkflowContext( context_id=context_id, workflow_name=kwargs.get( "workflow_name", f"workflow_{context_id[:8]}" ), workflow_steps=kwargs.get("workflow_steps", []), ) elif context_type == ContextType.SCENARIO: context = ScenarioContext( context_id=context_id, scenario_name=kwargs.get( "scenario_name", f"scenario_{context_id[:8]}" ), scenario_config=kwargs.get("scenario_config", {}), ) elif context_type == ContextType.PERFORMANCE: context = PerformanceContext( context_id=context_id, test_run_id=kwargs.get("test_run_id", f"run_{context_id[:8]}"), ) elif context_type == ContextType.AGENT: context = AgentContext( context_id=context_id, agent_id=kwargs.get("agent_id", f"agent_{context_id[:8]}"), agent_type=kwargs.get("agent_type", "custom"), agent_config=kwargs.get("agent_config", {}), ) else: context = BaseContext( context_id=context_id, context_type=context_type, data=kwargs.get("data", {}), parent_context_id=parent_context_id, metadata=kwargs.get("metadata", {}), ) # Set parent context if specified if parent_context_id: context.parent_context_id = parent_context_id # Store context self._contexts[context_id] = context self.storage.save_context(context) # Log context creation if self.audit_logger: self.audit_logger.log_context_operation( operation="create", context_id=context_id, context_type=context_type.value, details={"parent_context_id": parent_context_id}, ) return context def get_context(self, context_id: str) -> BaseContext | None: """Get context by ID, loading from storage if necessary.""" with self._lock: # Check in-memory cache first if context_id in self._contexts: return self._contexts[context_id] # Load from storage context = self.storage.load_context(context_id) if context: self._contexts[context_id] = context return context return None def update_context( self, context_id: str, updates: dict[str, Any], merge: bool = True ) -> bool: """Update context data.""" with self._lock: context = self.get_context(context_id) if not context: return False context.update_data(updates, merge) self.storage.save_context(context) # Log context update if self.audit_logger: self.audit_logger.log_context_operation( operation="update", context_id=context_id, context_type=context.context_type.value, details={"updates": updates, "merge": merge}, ) return True def delete_context(self, context_id: str) -> bool: """Delete context.""" with self._lock: # Remove from memory if context_id in self._contexts: context = self._contexts.pop(context_id) context_type = context.context_type.value else: context_type = "unknown" # Remove from storage deleted = self.storage.delete_context(context_id) # Log context deletion if deleted and self.audit_logger: self.audit_logger.log_context_operation( operation="delete", context_id=context_id, context_type=context_type ) return deleted def list_contexts( self, context_type: ContextType | None = None, status: ContextStatus | None = None, parent_context_id: str | None = None, tags: list[str] | None = None, ) -> list[dict[str, Any]]: """List contexts with optional filtering.""" return self.storage.list_contexts(context_type, status, parent_context_id, tags) def create_snapshot( self, context_id: str, description: str | None = None ) -> str | None: """Create a snapshot of context state.""" with self._lock: context = self.get_context(context_id) if not context: return None snapshot = context.create_snapshot(description) self.storage.save_context(context) # Log snapshot creation if self.audit_logger: self.audit_logger.log_context_operation( operation="snapshot", context_id=context_id, context_type=context.context_type.value, details={ "snapshot_id": snapshot.snapshot_id, "description": description, }, ) return snapshot.snapshot_id def restore_snapshot(self, context_id: str, snapshot_id: str) -> bool: """Restore context from snapshot.""" with self._lock: context = self.get_context(context_id) if not context: return False restored = context.restore_from_snapshot(snapshot_id) if restored: self.storage.save_context(context) # Log snapshot restoration if self.audit_logger: self.audit_logger.log_context_operation( operation="restore", context_id=context_id, context_type=context.context_type.value, details={"snapshot_id": snapshot_id}, ) return restored def get_global_context(self) -> GlobalContext: """Get the global context instance.""" return self._global_context def cleanup_expired_contexts(self) -> int: """Clean up expired contexts.""" with self._lock: # Remove expired contexts from memory expired_ids = [] for context_id, context in list(self._contexts.items()): if context.is_expired(): expired_ids.append(context_id) del self._contexts[context_id] # Clean up from storage storage_cleaned = self.storage.cleanup_expired_contexts() # Log cleanup if (expired_ids or storage_cleaned) and self.audit_logger: self.audit_logger.log_context_operation( operation="cleanup", context_id="system", context_type="cleanup", details={ "memory_cleaned": len(expired_ids), "storage_cleaned": storage_cleaned, }, ) return len(expired_ids) + storage_cleaned @asynccontextmanager async def context_session(self, context_type: ContextType, **kwargs): """Async context manager for automatic context lifecycle management.""" context = self.create_context(context_type, **kwargs) try: yield context finally: # Auto-save context on exit self.storage.save_context(context) def get_context_hierarchy(self, context_id: str) -> list[dict[str, Any]]: """Get the full hierarchy of a context (parent chain).""" hierarchy = [] current_id = context_id while current_id: context = self.get_context(current_id) if not context: break hierarchy.append( { "context_id": context.context_id, "context_type": context.context_type.value, "status": context.status.value, "created_at": context.created_at.isoformat(), "updated_at": context.updated_at.isoformat(), } ) current_id = context.parent_context_id return hierarchy def get_child_contexts(self, parent_context_id: str) -> list[dict[str, Any]]: """Get all child contexts of a parent context.""" return self.list_contexts(parent_context_id=parent_context_id) def shutdown(self) -> None: """Shutdown the context manager and save all contexts.""" with self._lock: # Save all in-memory contexts for context in self._contexts.values(): try: self.storage.save_context(context) except Exception: logger.exception(f"Failed to save context {context.context_id}") # Clear memory self._contexts.clear() # Log shutdown if self.audit_logger: self.audit_logger.log_context_operation( operation="shutdown", context_id="system", context_type="system" ) # Global context manager instance _global_context_manager: ContextManager | None = None def get_context_manager( db_path: str = "mcp_context.db", audit_logger: Any | None = None ) -> ContextManager: """Get or create the global context manager instance.""" global _global_context_manager # noqa: PLW0603 if _global_context_manager is None: _global_context_manager = ContextManager(db_path, audit_logger) return _global_context_manager def initialize_context_manager( db_path: str = "mcp_context.db", audit_logger: Any | None = None ) -> ContextManager: """Initialize the global context manager with specific configuration.""" global _global_context_manager # noqa: PLW0603 _global_context_manager = ContextManager(db_path, audit_logger) return _global_context_manager # Context management functions for MCP tools async def create_test_session_context( session_name: str, test_plan: dict[str, Any], session_config: dict[str, Any] | None = None, ) -> str: """Create a new test session context.""" manager = get_context_manager() context = manager.create_context( ContextType.TEST_SESSION, session_name=session_name, test_plan=test_plan, session_config=session_config, ) return context.context_id async def create_workflow_context( workflow_name: str, workflow_steps: list[dict[str, Any]], parent_context_id: str | None = None, ) -> str: """Create a new workflow context.""" manager = get_context_manager() context = manager.create_context( ContextType.WORKFLOW, workflow_name=workflow_name, workflow_steps=workflow_steps, parent_context_id=parent_context_id, ) return context.context_id async def create_agent_context( agent_id: str, agent_type: str, agent_config: dict[str, Any] | None = None ) -> str: """Create a new agent context.""" manager = get_context_manager() context = manager.create_context( ContextType.AGENT, agent_id=agent_id, agent_type=agent_type, agent_config=agent_config, ) return context.context_id async def get_context_data(context_id: str, key: str | None = None) -> Any | None: """Get data from a context.""" manager = get_context_manager() context = manager.get_context(context_id) if context: return context.get_data(key) return None async def update_context_data( context_id: str, updates: dict[str, Any], merge: bool = True ) -> bool: """Update context data.""" manager = get_context_manager() return manager.update_context(context_id, updates, merge) async def create_context_snapshot( context_id: str, description: str | None = None ) -> str | None: """Create a snapshot of context state.""" manager = get_context_manager() return manager.create_snapshot(context_id, description) async def restore_context_snapshot(context_id: str, snapshot_id: str) -> bool: """Restore context from snapshot.""" manager = get_context_manager() return manager.restore_snapshot(context_id, snapshot_id) async def list_contexts_by_type( context_type: ContextType, status: ContextStatus | None = None ) -> list[dict[str, Any]]: """List contexts by type and optional status.""" manager = get_context_manager() return manager.list_contexts(context_type=context_type, status=status) async def get_global_context_data(key: str | None = None) -> Any: """Get data from the global context.""" manager = get_context_manager() global_context = manager.get_global_context() return global_context.get_data(key) async def update_global_context_data( updates: dict[str, Any], merge: bool = True ) -> bool: """Update global context data.""" manager = get_context_manager() global_context = manager.get_global_context() global_context.update_data(updates, merge) manager.storage.save_context(global_context) return True