iflow-mcp_developermode-korea_reversecore-mcp 1.0.0__py3-none-any.whl

reversecore_mcp/core/resilience.py

@@ -0,0 +1,252 @@
+"""
+Resilience patterns for Reversecore_MCP.
+
+This module implements the Circuit Breaker pattern to prevent cascading failures
+when external tools (like Radare2 or Ghidra) become unstable.
+"""
+
+import functools
+import inspect
+import threading
+import time
+from collections.abc import Callable
+from enum import Enum
+from typing import TypeVar
+
+from reversecore_mcp.core.exceptions import ToolExecutionError
+from reversecore_mcp.core.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+F = TypeVar("F", bound=Callable)
+
+
+class CircuitState(Enum):
+    CLOSED = "CLOSED"  # Normal operation
+    OPEN = "OPEN"  # Failing, requests blocked
+    HALF_OPEN = "HALF_OPEN"  # Testing recovery
+
+
+class CircuitBreaker:
+    """
+    Circuit Breaker implementation with thread-safe state transitions.
+
+    If a tool fails 'failure_threshold' times within a window, the circuit opens
+    and blocks requests for 'recovery_timeout' seconds.
+    """
+
+    def __init__(self, name: str, failure_threshold: int = 5, recovery_timeout: int = 60):
+        self.name = name
+        self.failure_threshold = failure_threshold
+        self.recovery_timeout = recovery_timeout
+
+        self.state = CircuitState.CLOSED
+        self.failures = 0
+        self.last_failure_time = 0.0
+        self.next_attempt_time = 0.0
+        
+        # Thread lock for atomic state transitions
+        self._lock = threading.Lock()
+
+    def allow_request(self) -> bool:
+        """Check if a request is allowed (thread-safe)."""
+        with self._lock:
+            if self.state == CircuitState.CLOSED:
+                return True
+
+            if self.state == CircuitState.OPEN:
+                if time.time() >= self.next_attempt_time:
+                    logger.info(f"Circuit {self.name} entering HALF_OPEN state")
+                    self.state = CircuitState.HALF_OPEN
+                    return True
+                return False
+
+            if self.state == CircuitState.HALF_OPEN:
+                # Only one request allowed in HALF_OPEN - lock ensures this
+                return True
+
+            return True
+
+    def record_success(self):
+        """Record a successful execution (thread-safe)."""
+        with self._lock:
+            if self.state == CircuitState.HALF_OPEN:
+                logger.info(f"Circuit {self.name} recovered (CLOSED)")
+                self.state = CircuitState.CLOSED
+                self.failures = 0
+            elif self.state == CircuitState.CLOSED:
+                self.failures = 0
+
+    def record_failure(self):
+        """Record a failed execution (thread-safe)."""
+        with self._lock:
+            self.failures += 1
+            self.last_failure_time = time.time()
+
+            if self.state == CircuitState.CLOSED:
+                if self.failures >= self.failure_threshold:
+                    logger.warning(f"Circuit {self.name} opened due to {self.failures} failures")
+                    self.state = CircuitState.OPEN
+                    self.next_attempt_time = time.time() + self.recovery_timeout
+
+            elif self.state == CircuitState.HALF_OPEN:
+                logger.warning(f"Circuit {self.name} failed recovery, reopening")
+                self.state = CircuitState.OPEN
+                self.next_attempt_time = time.time() + self.recovery_timeout
+
+
+# Global registry of circuit breakers
+_breakers: dict[str, CircuitBreaker] = {}
+
+
+def get_circuit_breaker(name: str, **kwargs) -> CircuitBreaker:
+    """Get or create a circuit breaker for the given name."""
+    if name not in _breakers:
+        _breakers[name] = CircuitBreaker(name, **kwargs)
+    return _breakers[name]
+
+
+def circuit_breaker(
+    tool_name: str, failure_threshold: int = 5, recovery_timeout: int = 60
+) -> Callable[[F], F]:
+    """
+    Decorator to apply circuit breaker pattern to a function.
+
+    Automatically detects if the decorated function is async or sync
+    and applies the appropriate wrapper.
+
+    Args:
+        tool_name: Name of the tool for circuit breaker tracking
+        failure_threshold: Number of failures before opening circuit
+        recovery_timeout: Seconds to wait before attempting recovery
+
+    Returns:
+        Decorated function with circuit breaker protection
+    """
+
+    def decorator(func: F) -> F:
+        breaker = get_circuit_breaker(
+            tool_name,
+            failure_threshold=failure_threshold,
+            recovery_timeout=recovery_timeout,
+        )
+
+        def _get_error_message() -> str:
+            """Generate error message for circuit open state."""
+            remaining = int(breaker.next_attempt_time - time.time())
+            return (
+                f"Tool '{tool_name}' is temporarily unavailable due to repeated failures. "
+                f"Please try again in {max(0, remaining)} seconds."
+            )
+
+        # Check if function is async
+        if inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                if not breaker.allow_request():
+                    raise ToolExecutionError(_get_error_message())
+
+                try:
+                    result = await func(*args, **kwargs)
+                    breaker.record_success()
+                    return result
+                except Exception:
+                    breaker.record_failure()
+                    raise
+
+            return async_wrapper  # type: ignore[return-value]
+
+        # Sync version
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            if not breaker.allow_request():
+                raise ToolExecutionError(_get_error_message())
+
+            try:
+                result = func(*args, **kwargs)
+                breaker.record_success()
+                return result
+            except Exception:
+                breaker.record_failure()
+                raise
+
+        return sync_wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def circuit_breaker_sync(
+    tool_name: str, failure_threshold: int = 5, recovery_timeout: int = 60
+) -> Callable[[F], F]:
+    """
+    Explicit sync-only circuit breaker decorator.
+
+    Use this when you want to explicitly mark a function as sync,
+    or when the auto-detection in circuit_breaker doesn't work correctly.
+    """
+
+    def decorator(func: F) -> F:
+        breaker = get_circuit_breaker(
+            tool_name,
+            failure_threshold=failure_threshold,
+            recovery_timeout=recovery_timeout,
+        )
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            if not breaker.allow_request():
+                raise ToolExecutionError(
+                    f"Tool '{tool_name}' is temporarily unavailable due to repeated failures. "
+                    f"Please try again in {int(breaker.next_attempt_time - time.time())} seconds."
+                )
+
+            try:
+                result = func(*args, **kwargs)
+                breaker.record_success()
+                return result
+            except Exception:
+                breaker.record_failure()
+                raise
+
+        return wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def circuit_breaker_async(
+    tool_name: str, failure_threshold: int = 5, recovery_timeout: int = 60
+) -> Callable[[F], F]:
+    """
+    Explicit async-only circuit breaker decorator.
+
+    Use this when you want to explicitly mark a function as async,
+    or when the auto-detection in circuit_breaker doesn't work correctly.
+    """
+
+    def decorator(func: F) -> F:
+        breaker = get_circuit_breaker(
+            tool_name,
+            failure_threshold=failure_threshold,
+            recovery_timeout=recovery_timeout,
+        )
+
+        @functools.wraps(func)
+        async def wrapper(*args, **kwargs):
+            if not breaker.allow_request():
+                raise ToolExecutionError(
+                    f"Tool '{tool_name}' is temporarily unavailable due to repeated failures. "
+                    f"Please try again in {int(breaker.next_attempt_time - time.time())} seconds."
+                )
+
+            try:
+                result = await func(*args, **kwargs)
+                breaker.record_success()
+                return result
+            except Exception:
+                breaker.record_failure()
+                raise
+
+        return wrapper  # type: ignore[return-value]
+
+    return decorator

reversecore_mcp/core/resource_manager.py

@@ -0,0 +1,169 @@
+"""
+Resource Manager for Reversecore_MCP.
+
+This module handles periodic cleanup of temporary files and stale cache entries
+to prevent resource exhaustion over time.
+"""
+
+import asyncio
+import time
+
+import os
+from contextlib import suppress
+
+from reversecore_mcp.core import config
+from reversecore_mcp.core.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+class ResourceManager:
+    """
+    Manages background cleanup tasks.
+    
+    Responsibilities:
+    1. Clean up stale temporary files (disk)
+    2. Reap zombie processes (kernel)
+    """
+
+    def __init__(self, cleanup_interval: int = 3600, pid_check_interval: int = 60):  # Default: 1h files, 60s PIDs
+        self.cleanup_interval = cleanup_interval
+        self.pid_check_interval = pid_check_interval
+        self._task: asyncio.Task | None = None
+        self._pid_task: asyncio.Task | None = None
+        self._running = False
+        self._tracked_pids: set[int] = set()
+
+    def track_pid(self, pid: int) -> None:
+        """Track a subprocess PID for zombie cleanup."""
+        self._tracked_pids.add(pid)
+
+    async def start(self):
+        """Start the background cleanup tasks."""
+        if self._running:
+            return
+
+        self._running = True
+        self._task = asyncio.create_task(self._cleanup_loop())
+        self._pid_task = asyncio.create_task(self._pid_check_loop())
+        logger.info(
+            f"Resource Manager started (cleanup: {self.cleanup_interval}s, pid_check: {self.pid_check_interval}s)"
+        )
+
+    async def stop(self):
+        """Stop all background cleanup tasks."""
+        self._running = False
+        for task in [self._task, self._pid_task]:
+            if task:
+                task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await task
+        logger.info("Resource Manager stopped")
+
+    async def _cleanup_loop(self):
+        """Main resource cleanup loop (files, logs)."""
+        while self._running:
+            try:
+                await asyncio.sleep(self.cleanup_interval)
+                await self.cleanup()
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Error in cleanup loop: {e}")
+
+    async def _pid_check_loop(self):
+        """FAST PID health check loop (zombies)."""
+        while self._running:
+            try:
+                await asyncio.sleep(self.pid_check_interval)
+                self._reap_zombies()
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Error in PID check loop: {e}")
+
+    def _reap_zombies(self):
+        """Reap tracked zombie processes."""
+        if not self._tracked_pids:
+            return
+
+        reaped_count = 0
+        dead_pids = set()
+
+        for pid in list(self._tracked_pids):
+            try:
+                # Check if process is still alive and defunct
+                # waitpid with WNOHANG returns (pid, status) if dead, (0, 0) if running
+                wpid, status = os.waitpid(pid, os.WNOHANG)
+                if wpid > 0:
+                    # Process was a zombie and is now reaped
+                    dead_pids.add(pid)
+                    reaped_count += 1
+            except ChildProcessError:
+                # Process already gone or not our child
+                dead_pids.add(pid)
+            except Exception as e:
+                logger.debug(f"Failed to check PID {pid}: {e}")
+
+        if dead_pids:
+            self._tracked_pids -= dead_pids
+            if reaped_count > 0:
+                logger.info(f"Reaped {reaped_count} zombie processes")
+
+    async def cleanup(self):
+        """Perform cleanup operations."""
+        logger.info("Starting periodic resource cleanup...")
+
+        # 1. Clear stale binary cache
+        # (BinaryMetadataCache already checks mtime, but we can clear old entries if needed)
+        # For now, we just log.
+        # binary_cache.clear() # Too aggressive?
+
+        # 2. Clean up temporary files
+        cfg = config.get_config()
+        workspace = cfg.workspace
+
+        try:
+            # Clean .tmp files older than 24 hours
+            now = time.time()
+            max_age = 24 * 3600
+
+            cleaned_count = 0
+
+            # OPTIMIZATION: Use itertools.chain to avoid multiple glob calls and iterations
+            # This combines all temp file patterns into a single iterable for better performance
+            from itertools import chain
+
+            # Combine all patterns into a single iterable
+            temp_files = chain(
+                workspace.glob("*.tmp"), workspace.glob(".r2_*"), workspace.glob("*.r2")
+            )
+
+            # PERFORMANCE NOTE: For very large numbers of temp files (>1000),
+            # consider using batch deletion with os.unlink_many() or parallel deletion
+            # However, this is a rare case and the current implementation is sufficient
+            for temp_file in temp_files:
+                try:
+                    if temp_file.is_file():
+                        mtime = temp_file.stat().st_mtime
+                        if now - mtime > max_age:
+                            temp_file.unlink()
+                            cleaned_count += 1
+                except Exception as e:
+                    logger.warning(f"Failed to delete temp file {temp_file}: {e}")
+
+            if cleaned_count > 0:
+                logger.info(f"Cleaned up {cleaned_count} stale temporary files")
+
+        except Exception as e:
+            logger.error(f"Failed to clean temp files: {e}")
+
+        # 3. Check r2_pool health?
+        # r2_pool manages itself via LRU.
+
+        logger.info("Resource cleanup complete")
+
+
+# Global instance (for backward compatibility)
+# New code should use: from reversecore_mcp.core.container import get_resource_manager
+resource_manager = ResourceManager()

reversecore_mcp/core/result.py

@@ -0,0 +1,132 @@
+"""Pydantic models for structured tool results."""
+
+from __future__ import annotations
+
+from typing import Any, Literal, Union
+
+from pydantic import BaseModel
+
+try:
+    from typing import NotRequired, TypedDict
+except ImportError:
+    from typing_extensions import NotRequired, TypedDict
+
+
+# TypedDict definitions for common tool result structures
+class FunctionInfo(TypedDict):
+    """Information about a function in a binary."""
+
+    name: str
+    address: str
+    size: NotRequired[int]
+    signature: NotRequired[str]
+    callees: NotRequired[list[str]]
+    callers: NotRequired[list[str]]
+
+
+class DisassemblyResult(TypedDict):
+    """Result of disassembly operation."""
+
+    address: str
+    mnemonic: str
+    operands: str
+    bytes: NotRequired[str]
+    comment: NotRequired[str]
+
+
+class DecompilationResult(TypedDict):
+    """Result of decompilation operation."""
+
+    function_name: str
+    source_code: str
+    decompiler: NotRequired[str]
+    address: NotRequired[str]
+
+
+class BinaryMetadata(TypedDict):
+    """Metadata about a binary file."""
+
+    file_path: str
+    file_size: int
+    file_type: str
+    architecture: NotRequired[str]
+    endianness: NotRequired[str]
+    entry_point: NotRequired[str]
+    sections: NotRequired[list[dict[str, Any]]]
+
+
+class YaraRuleResult(TypedDict):
+    """Result of YARA rule generation."""
+
+    rule_name: str
+    rule_content: str
+    patterns_count: NotRequired[int]
+    meta: NotRequired[dict[str, str]]
+
+
+class ScanResult(TypedDict):
+    """Result of a security scan."""
+
+    findings: list[dict[str, Any]]
+    severity: NotRequired[str]
+    recommendations: NotRequired[list[str]]
+
+
+class EmulationResult(TypedDict):
+    """Result of code emulation."""
+
+    final_registers: dict[str, Any]
+    steps_executed: int
+    status: str
+    memory_writes: NotRequired[list[dict[str, Any]]]
+    syscalls: NotRequired[list[str]]
+
+
+class ErrorDetails(TypedDict, total=False):
+    """Details for error responses."""
+
+    max_size: int
+    actual_size: int
+    exception_type: str
+    timeout_seconds: int
+
+
+class ToolSuccess(BaseModel):
+    """Represents a successful tool invocation."""
+
+    status: Literal["success"] = "success"
+    data: str | dict[str, Any]
+    metadata: dict[str, Any] | None = None
+
+
+class ToolError(BaseModel):
+    """Represents a failed tool invocation."""
+
+    status: Literal["error"] = "error"
+    error_code: str
+    message: str
+    hint: str | None = None
+    details: dict[str, Any] | None = None
+
+
+ToolResult = Union[ToolSuccess, ToolError]
+
+
+def success(data: str | dict[str, Any], **metadata: Any) -> ToolSuccess:
+    """Create a ToolSuccess instance with optional metadata."""
+    return ToolSuccess(data=data, metadata=metadata or None)
+
+
+def failure(
+    error_code: str,
+    message: str,
+    hint: str | None = None,
+    **details: Any,
+) -> ToolError:
+    """Create a ToolError instance with optional hint/details."""
+    return ToolError(
+        error_code=error_code,
+        message=message,
+        hint=hint,
+        details=details or None,
+    )