mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-26 11:49:36 +00:00
feat: Add universal auto-timeout and backgrounding for MCP tools
- Implement MCPToolTimeoutManager with configurable timeouts - Add @auto_timeout decorator for all MCP tools - Apply auto-timeout to search, find, and ast tools - Default 2-minute timeout, configurable via HANZO_MCP_TOOL_TIMEOUT - Background operations continue with full log files - Process management integration for monitoring and termination - Solves MCP error -32001 (Request timed out) issues Version bump: 0.8.15 → 0.8.16
This commit is contained in:
@@ -0,0 +1,129 @@
|
||||
# Security Patch for Hanzo MCP v0.8.13
|
||||
|
||||
## Critical Security Vulnerabilities Found
|
||||
|
||||
### 1. Command Injection (CRITICAL) 🚨
|
||||
**Location**: `hanzo_mcp/tools/shell/command_executor.py` line 114
|
||||
**Issue**: Insufficient shell escaping using basic string replacement
|
||||
**Risk**: Remote code execution
|
||||
|
||||
**Current Code**:
|
||||
```python
|
||||
escaped_command = command.replace('"', '\\"')
|
||||
```
|
||||
|
||||
**Fixed Code**:
|
||||
```python
|
||||
import shlex
|
||||
escaped_command = shlex.quote(command)
|
||||
```
|
||||
|
||||
### 2. Path Traversal (HIGH) ⚠️
|
||||
**Location**: `hanzo_mcp/tools/common/permissions.py`
|
||||
**Issue**: No validation for symlinks and `../` traversal
|
||||
**Risk**: Unauthorized file access
|
||||
|
||||
**Add validation**:
|
||||
```python
|
||||
def validate_path(self, path: str) -> bool:
|
||||
"""Validate path against traversal attacks."""
|
||||
resolved = Path(path).resolve()
|
||||
|
||||
# Check for path traversal
|
||||
if ".." in str(path):
|
||||
return False
|
||||
|
||||
# Check against allowed paths
|
||||
for allowed in self.allowed_paths:
|
||||
try:
|
||||
resolved.relative_to(allowed)
|
||||
return True
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
return False
|
||||
```
|
||||
|
||||
### 3. Missing Authentication (HIGH) ⚠️
|
||||
**Location**: `hanzo_mcp/server.py`
|
||||
**Issue**: MCP server accepts any local connection without auth
|
||||
**Risk**: Unauthorized access to system
|
||||
|
||||
**Add basic token auth**:
|
||||
```python
|
||||
import secrets
|
||||
import os
|
||||
|
||||
class HanzoMCPServer:
|
||||
def __init__(self, ...):
|
||||
self.auth_token = os.environ.get('HANZO_MCP_TOKEN') or secrets.token_urlsafe(32)
|
||||
|
||||
def verify_auth(self, request_token: str) -> bool:
|
||||
"""Verify authentication token."""
|
||||
return secrets.compare_digest(request_token, self.auth_token)
|
||||
```
|
||||
|
||||
### 4. Subprocess Timeout (MEDIUM)
|
||||
**Location**: Multiple subprocess.run() calls
|
||||
**Issue**: Missing consistent timeout enforcement
|
||||
**Risk**: Resource exhaustion from hanging processes
|
||||
|
||||
**Fix**:
|
||||
```python
|
||||
# Always use timeout
|
||||
result = subprocess.run(
|
||||
command,
|
||||
timeout=self.command_timeout, # Default 120s
|
||||
check=False,
|
||||
capture_output=True
|
||||
)
|
||||
```
|
||||
|
||||
## Test Results Summary
|
||||
|
||||
- **Total Tests**: 410 collected
|
||||
- **Test Errors**: 4 collection errors
|
||||
- **Status**: Tests not fully passing due to environment issues
|
||||
- **Recommendation**: Fix security issues first, then resolve test environment
|
||||
|
||||
## Immediate Actions Required
|
||||
|
||||
1. **Apply command injection patch** using shlex.quote()
|
||||
2. **Add path validation** to PermissionManager
|
||||
3. **Implement token authentication** for MCP server
|
||||
4. **Add timeout to all subprocess calls**
|
||||
5. **Add audit logging** for all operations
|
||||
6. **Implement rate limiting** to prevent abuse
|
||||
|
||||
## Risk Assessment
|
||||
|
||||
| Vulnerability | Severity | Exploitability | Impact | Priority |
|
||||
|--------------|----------|----------------|---------|----------|
|
||||
| Command Injection | CRITICAL | High | RCE | P0 |
|
||||
| Path Traversal | HIGH | Medium | File Access | P1 |
|
||||
| No Authentication | HIGH | High | Unauthorized Access | P1 |
|
||||
| Missing Timeouts | MEDIUM | Low | DoS | P2 |
|
||||
|
||||
## Deployment Recommendation
|
||||
|
||||
**DO NOT DEPLOY TO PRODUCTION** until:
|
||||
1. All CRITICAL and HIGH severity issues are patched
|
||||
2. Authentication is implemented
|
||||
3. Test suite is passing (currently has environment issues)
|
||||
4. Security audit is performed on patched version
|
||||
|
||||
## Next Version Plan (v0.8.14)
|
||||
|
||||
Should include:
|
||||
- Security patches for all identified vulnerabilities
|
||||
- Authentication system
|
||||
- Audit logging
|
||||
- Rate limiting
|
||||
- Improved error handling
|
||||
- Complete test coverage
|
||||
|
||||
---
|
||||
|
||||
Generated by Code Review Agent
|
||||
Date: 2024-12-20
|
||||
Version Reviewed: hanzo-mcp v0.8.13
|
||||
@@ -1,31 +0,0 @@
|
||||
"""Hanzo AI - Implementation of Hanzo capabilities using MCP."""
|
||||
|
||||
# Polyfill typing.override for Python < 3.12
|
||||
try: # pragma: no cover
|
||||
from typing import override as _override # type: ignore
|
||||
except Exception: # pragma: no cover
|
||||
import typing as _typing
|
||||
|
||||
def override(obj): # type: ignore
|
||||
return obj
|
||||
|
||||
_typing.override = override # type: ignore[attr-defined]
|
||||
|
||||
# Configure FastMCP logging globally for stdio transport
|
||||
import os
|
||||
import warnings
|
||||
|
||||
# Suppress litellm deprecation warnings about event loop
|
||||
warnings.filterwarnings(
|
||||
"ignore", message="There is no current event loop", category=DeprecationWarning
|
||||
)
|
||||
|
||||
if os.environ.get("HANZO_MCP_TRANSPORT") == "stdio":
|
||||
try:
|
||||
from fastmcp.utilities.logging import configure_logging
|
||||
|
||||
configure_logging(level="ERROR")
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
__version__ = "0.8.4"
|
||||
@@ -1,6 +0,0 @@
|
||||
"""Main entry point for hanzo-mcp when run as a module."""
|
||||
|
||||
from hanzo_mcp.cli import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,5 +0,0 @@
|
||||
"""Analytics module for Hanzo MCP using PostHog."""
|
||||
|
||||
from .posthog_analytics import Analytics, track_error, track_event, track_tool_usage
|
||||
|
||||
__all__ = ["Analytics", "track_event", "track_tool_usage", "track_error"]
|
||||
@@ -1,358 +0,0 @@
|
||||
"""PostHog analytics integration for Hanzo MCP.
|
||||
|
||||
This module provides analytics tracking for:
|
||||
- Tool usage and performance
|
||||
- Error tracking and debugging
|
||||
- Feature adoption and user behavior
|
||||
- A/B testing and feature flags
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import asyncio
|
||||
import platform
|
||||
import functools
|
||||
import traceback
|
||||
from typing import Any, Dict, TypeVar, Callable, Optional
|
||||
from datetime import datetime
|
||||
from dataclasses import dataclass
|
||||
|
||||
# Try to import PostHog, but make it optional
|
||||
try:
|
||||
from posthog import Posthog
|
||||
|
||||
POSTHOG_AVAILABLE = True
|
||||
except ImportError:
|
||||
POSTHOG_AVAILABLE = False
|
||||
Posthog = None
|
||||
|
||||
|
||||
F = TypeVar("F", bound=Callable[..., Any])
|
||||
|
||||
|
||||
@dataclass
|
||||
class AnalyticsConfig:
|
||||
"""Configuration for analytics."""
|
||||
|
||||
api_key: Optional[str] = None
|
||||
host: str = "https://us.i.posthog.com"
|
||||
enabled: bool = True
|
||||
debug: bool = False
|
||||
capture_errors: bool = True
|
||||
capture_performance: bool = True
|
||||
distinct_id: Optional[str] = None
|
||||
|
||||
|
||||
class Analytics:
|
||||
"""Main analytics class for Hanzo MCP."""
|
||||
|
||||
def __init__(self, config: Optional[AnalyticsConfig] = None):
|
||||
"""Initialize analytics with configuration."""
|
||||
self.config = config or AnalyticsConfig()
|
||||
self._client = None
|
||||
|
||||
# Load from environment if not provided
|
||||
if not self.config.api_key:
|
||||
self.config.api_key = os.environ.get("POSTHOG_API_KEY")
|
||||
|
||||
if not self.config.distinct_id:
|
||||
# Use machine ID or generate one
|
||||
self.config.distinct_id = self._get_distinct_id()
|
||||
|
||||
# Initialize PostHog if available and configured
|
||||
if POSTHOG_AVAILABLE and self.config.api_key and self.config.enabled:
|
||||
self._client = Posthog(
|
||||
self.config.api_key,
|
||||
host=self.config.host,
|
||||
debug=self.config.debug,
|
||||
enable_exception_autocapture=self.config.capture_errors,
|
||||
)
|
||||
|
||||
def _get_distinct_id(self) -> str:
|
||||
"""Get a distinct ID for this installation."""
|
||||
# Try to get from environment
|
||||
distinct_id = os.environ.get("HANZO_DISTINCT_ID")
|
||||
if distinct_id:
|
||||
return distinct_id
|
||||
|
||||
# Use hostname + username as fallback
|
||||
import socket
|
||||
import getpass
|
||||
|
||||
hostname = socket.gethostname()
|
||||
username = getpass.getuser()
|
||||
return f"{hostname}:{username}"
|
||||
|
||||
def is_enabled(self) -> bool:
|
||||
"""Check if analytics is enabled."""
|
||||
return bool(self._client and self.config.enabled)
|
||||
|
||||
def capture(self, event: str, properties: Optional[Dict[str, Any]] = None) -> None:
|
||||
"""Capture an analytics event."""
|
||||
if not self.is_enabled():
|
||||
return
|
||||
|
||||
try:
|
||||
# Add common properties
|
||||
props = {
|
||||
"timestamp": datetime.utcnow().isoformat(),
|
||||
"platform": platform.system(),
|
||||
"python_version": platform.python_version(),
|
||||
"mcp_version": "0.6.13", # TODO: Get from package
|
||||
**(properties or {}),
|
||||
}
|
||||
|
||||
self._client.capture(self.config.distinct_id, event, properties=props)
|
||||
except Exception as e:
|
||||
if self.config.debug:
|
||||
print(f"Analytics error: {e}")
|
||||
|
||||
def identify(self, properties: Optional[Dict[str, Any]] = None) -> None:
|
||||
"""Identify the current user/installation."""
|
||||
if not self.is_enabled():
|
||||
return
|
||||
|
||||
try:
|
||||
self._client.identify(self.config.distinct_id, properties=properties or {})
|
||||
except Exception as e:
|
||||
if self.config.debug:
|
||||
print(f"Analytics identify error: {e}")
|
||||
|
||||
def track_tool_usage(
|
||||
self,
|
||||
tool_name: str,
|
||||
duration_ms: Optional[float] = None,
|
||||
success: bool = True,
|
||||
error: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Track tool usage event."""
|
||||
properties = {"tool_name": tool_name, "success": success, **(metadata or {})}
|
||||
|
||||
if duration_ms is not None:
|
||||
properties["duration_ms"] = duration_ms
|
||||
|
||||
if error:
|
||||
properties["error"] = str(error)
|
||||
|
||||
self.capture("tool_used", properties)
|
||||
|
||||
def track_error(
|
||||
self, error: Exception, context: Optional[Dict[str, Any]] = None
|
||||
) -> None:
|
||||
"""Track an error event."""
|
||||
if not self.config.capture_errors:
|
||||
return
|
||||
|
||||
properties = {
|
||||
"error_type": type(error).__name__,
|
||||
"error_message": str(error),
|
||||
"error_traceback": traceback.format_exc(),
|
||||
**(context or {}),
|
||||
}
|
||||
|
||||
self.capture("error_occurred", properties)
|
||||
|
||||
def feature_enabled(self, flag_key: str, default: bool = False) -> bool:
|
||||
"""Check if a feature flag is enabled."""
|
||||
if not self.is_enabled():
|
||||
return default
|
||||
|
||||
try:
|
||||
return self._client.feature_enabled(
|
||||
flag_key, self.config.distinct_id, default=default
|
||||
)
|
||||
except Exception:
|
||||
return default
|
||||
|
||||
def get_feature_flag(self, flag_key: str, default: Any = None) -> Any:
|
||||
"""Get feature flag value."""
|
||||
if not self.is_enabled():
|
||||
return default
|
||||
|
||||
try:
|
||||
return self._client.get_feature_flag(
|
||||
flag_key, self.config.distinct_id, default=default
|
||||
)
|
||||
except Exception:
|
||||
return default
|
||||
|
||||
def flush(self) -> None:
|
||||
"""Flush any pending events."""
|
||||
if self.is_enabled():
|
||||
try:
|
||||
self._client.flush()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Shutdown analytics client."""
|
||||
if self.is_enabled():
|
||||
try:
|
||||
self._client.shutdown()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# Global analytics instance
|
||||
_analytics = None
|
||||
|
||||
|
||||
def get_analytics() -> Analytics:
|
||||
"""Get or create the global analytics instance."""
|
||||
global _analytics
|
||||
if _analytics is None:
|
||||
_analytics = Analytics()
|
||||
return _analytics
|
||||
|
||||
|
||||
def track_event(event: str, properties: Optional[Dict[str, Any]] = None) -> None:
|
||||
"""Track a custom event."""
|
||||
get_analytics().capture(event, properties)
|
||||
|
||||
|
||||
def track_tool_usage(tool_name: str, **kwargs) -> None:
|
||||
"""Track tool usage."""
|
||||
get_analytics().track_tool_usage(tool_name, **kwargs)
|
||||
|
||||
|
||||
def track_error(error: Exception, context: Optional[Dict[str, Any]] = None) -> None:
|
||||
"""Track an error."""
|
||||
get_analytics().track_error(error, context)
|
||||
|
||||
|
||||
def with_analytics(tool_name: str):
|
||||
"""Decorator to track tool usage with analytics."""
|
||||
|
||||
def decorator(func: F) -> F:
|
||||
@functools.wraps(func)
|
||||
async def async_wrapper(*args, **kwargs):
|
||||
start_time = time.time()
|
||||
error = None
|
||||
try:
|
||||
result = await func(*args, **kwargs)
|
||||
return result
|
||||
except Exception as e:
|
||||
error = e
|
||||
track_error(e, {"tool": tool_name})
|
||||
raise
|
||||
finally:
|
||||
duration_ms = (time.time() - start_time) * 1000
|
||||
track_tool_usage(
|
||||
tool_name,
|
||||
duration_ms=duration_ms,
|
||||
success=error is None,
|
||||
error=str(error) if error else None,
|
||||
)
|
||||
|
||||
@functools.wraps(func)
|
||||
def sync_wrapper(*args, **kwargs):
|
||||
start_time = time.time()
|
||||
error = None
|
||||
try:
|
||||
result = func(*args, **kwargs)
|
||||
return result
|
||||
except Exception as e:
|
||||
error = e
|
||||
track_error(e, {"tool": tool_name})
|
||||
raise
|
||||
finally:
|
||||
duration_ms = (time.time() - start_time) * 1000
|
||||
track_tool_usage(
|
||||
tool_name,
|
||||
duration_ms=duration_ms,
|
||||
success=error is None,
|
||||
error=str(error) if error else None,
|
||||
)
|
||||
|
||||
# Return appropriate wrapper based on function type
|
||||
if asyncio.iscoroutinefunction(func):
|
||||
return async_wrapper
|
||||
else:
|
||||
return sync_wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def feature_flag(flag_key: str, default: bool = False):
|
||||
"""Decorator to conditionally enable features based on flags."""
|
||||
|
||||
def decorator(func: F) -> F:
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
if get_analytics().feature_enabled(flag_key, default):
|
||||
return func(*args, **kwargs)
|
||||
else:
|
||||
raise NotImplementedError(f"Feature '{flag_key}' is not enabled")
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
# Tool usage context manager
|
||||
class ToolUsageTracker:
|
||||
"""Context manager for tracking tool usage."""
|
||||
|
||||
def __init__(self, tool_name: str, metadata: Optional[Dict[str, Any]] = None):
|
||||
self.tool_name = tool_name
|
||||
self.metadata = metadata or {}
|
||||
self.start_time = None
|
||||
self.error = None
|
||||
|
||||
def __enter__(self):
|
||||
self.start_time = time.time()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
duration_ms = (time.time() - self.start_time) * 1000
|
||||
success = exc_type is None
|
||||
|
||||
if exc_type:
|
||||
self.error = str(exc_val)
|
||||
track_error(exc_val, {"tool": self.tool_name, **self.metadata})
|
||||
|
||||
track_tool_usage(
|
||||
self.tool_name,
|
||||
duration_ms=duration_ms,
|
||||
success=success,
|
||||
error=self.error,
|
||||
metadata=self.metadata,
|
||||
)
|
||||
|
||||
# Don't suppress exceptions
|
||||
return False
|
||||
|
||||
|
||||
# Async context manager version
|
||||
class AsyncToolUsageTracker:
|
||||
"""Async context manager for tracking tool usage."""
|
||||
|
||||
def __init__(self, tool_name: str, metadata: Optional[Dict[str, Any]] = None):
|
||||
self.tool_name = tool_name
|
||||
self.metadata = metadata or {}
|
||||
self.start_time = None
|
||||
self.error = None
|
||||
|
||||
async def __aenter__(self):
|
||||
self.start_time = time.time()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
duration_ms = (time.time() - self.start_time) * 1000
|
||||
success = exc_type is None
|
||||
|
||||
if exc_type:
|
||||
self.error = str(exc_val)
|
||||
track_error(exc_val, {"tool": self.tool_name, **self.metadata})
|
||||
|
||||
track_tool_usage(
|
||||
self.tool_name,
|
||||
duration_ms=duration_ms,
|
||||
success=success,
|
||||
error=self.error,
|
||||
metadata=self.metadata,
|
||||
)
|
||||
|
||||
# Don't suppress exceptions
|
||||
return False
|
||||
@@ -1,484 +0,0 @@
|
||||
"""MCP Bridge for inter-Claude communication.
|
||||
|
||||
This module provides MCP server functionality that allows Claude instances
|
||||
to communicate with each other, enabling peer-to-peer agent networks.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import asyncio
|
||||
import logging
|
||||
import argparse
|
||||
from typing import Any, Dict, List, Optional
|
||||
from dataclasses import asdict, dataclass
|
||||
|
||||
import mcp.server.fastmcp as mcp
|
||||
from mcp import tool
|
||||
from mcp.types import INTERNAL_ERROR, Tool, TextContent
|
||||
from mcp.server.stdio import stdio_server
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BridgeConfig:
|
||||
"""Configuration for MCP bridge."""
|
||||
|
||||
target_port: int
|
||||
instance_id: int
|
||||
role: str
|
||||
source_instance: Optional[int] = None
|
||||
target_instance: Optional[int] = None
|
||||
|
||||
|
||||
class ClaudeBridge(FastMCP):
|
||||
"""MCP Bridge server for Claude-to-Claude communication."""
|
||||
|
||||
def __init__(self, config: BridgeConfig):
|
||||
"""Initialize the bridge.
|
||||
|
||||
Args:
|
||||
config: Bridge configuration
|
||||
"""
|
||||
# Set server name based on target instance
|
||||
super().__init__(f"claude_instance_{config.instance_id}")
|
||||
|
||||
self.config = config
|
||||
self.conversation_history: List[Dict[str, Any]] = []
|
||||
self.shared_context: Dict[str, Any] = {}
|
||||
|
||||
# Register tools
|
||||
self._register_tools()
|
||||
|
||||
def _register_tools(self):
|
||||
"""Register MCP tools for inter-Claude communication."""
|
||||
|
||||
@self.tool()
|
||||
async def chat_with_claude(message: str, context: Optional[str] = None) -> str:
|
||||
"""Chat with another Claude instance.
|
||||
|
||||
Args:
|
||||
message: Message to send to the other Claude
|
||||
context: Optional context to provide
|
||||
|
||||
Returns:
|
||||
Response from the other Claude instance
|
||||
"""
|
||||
logger.info(f"Bridge {self.config.instance_id}: Received chat request")
|
||||
|
||||
# Record in conversation history
|
||||
self.conversation_history.append(
|
||||
{
|
||||
"from": self.config.source_instance,
|
||||
"to": self.config.target_instance,
|
||||
"message": message,
|
||||
"context": context,
|
||||
}
|
||||
)
|
||||
|
||||
# Simulate response (in production, this would make actual API call)
|
||||
response = await self._forward_to_claude(message, context)
|
||||
|
||||
self.conversation_history.append(
|
||||
{
|
||||
"from": self.config.target_instance,
|
||||
"to": self.config.source_instance,
|
||||
"response": response,
|
||||
}
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
@self.tool()
|
||||
async def ask_claude_to_review(
|
||||
code: str, description: str, focus_areas: Optional[List[str]] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""Ask another Claude to review code.
|
||||
|
||||
Args:
|
||||
code: Code to review
|
||||
description: Description of what the code does
|
||||
focus_areas: Specific areas to focus on (e.g., ["security", "performance"])
|
||||
|
||||
Returns:
|
||||
Review feedback from the other Claude
|
||||
"""
|
||||
logger.info(f"Bridge {self.config.instance_id}: Code review request")
|
||||
|
||||
review_prompt = self._build_review_prompt(code, description, focus_areas)
|
||||
review = await self._forward_to_claude(review_prompt)
|
||||
|
||||
return {
|
||||
"reviewer": f"claude_{self.config.instance_id}",
|
||||
"role": self.config.role,
|
||||
"feedback": review,
|
||||
"focus_areas": focus_areas or ["general"],
|
||||
}
|
||||
|
||||
@self.tool()
|
||||
async def delegate_to_claude(
|
||||
task: str, requirements: List[str], constraints: Optional[List[str]] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""Delegate a task to another Claude instance.
|
||||
|
||||
Args:
|
||||
task: Task description
|
||||
requirements: List of requirements
|
||||
constraints: Optional constraints
|
||||
|
||||
Returns:
|
||||
Task completion result from the other Claude
|
||||
"""
|
||||
logger.info(f"Bridge {self.config.instance_id}: Task delegation")
|
||||
|
||||
delegation_prompt = self._build_delegation_prompt(
|
||||
task, requirements, constraints
|
||||
)
|
||||
result = await self._forward_to_claude(delegation_prompt)
|
||||
|
||||
return {
|
||||
"delegated_to": f"claude_{self.config.instance_id}",
|
||||
"role": self.config.role,
|
||||
"task": task,
|
||||
"result": result,
|
||||
"status": "completed",
|
||||
}
|
||||
|
||||
@self.tool()
|
||||
async def get_claude_opinion(
|
||||
question: str,
|
||||
options: Optional[List[str]] = None,
|
||||
criteria: Optional[List[str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Get another Claude's opinion on a decision.
|
||||
|
||||
Args:
|
||||
question: The question or decision to get opinion on
|
||||
options: Optional list of options to choose from
|
||||
criteria: Optional evaluation criteria
|
||||
|
||||
Returns:
|
||||
Opinion and reasoning from the other Claude
|
||||
"""
|
||||
logger.info(f"Bridge {self.config.instance_id}: Opinion request")
|
||||
|
||||
opinion_prompt = self._build_opinion_prompt(question, options, criteria)
|
||||
opinion = await self._forward_to_claude(opinion_prompt)
|
||||
|
||||
return {
|
||||
"advisor": f"claude_{self.config.instance_id}",
|
||||
"role": self.config.role,
|
||||
"question": question,
|
||||
"opinion": opinion,
|
||||
"options_considered": options,
|
||||
"criteria_used": criteria,
|
||||
}
|
||||
|
||||
@self.tool()
|
||||
async def share_context_with_claude(
|
||||
key: str, value: Any, description: Optional[str] = None
|
||||
) -> bool:
|
||||
"""Share context with another Claude instance.
|
||||
|
||||
Args:
|
||||
key: Context key
|
||||
value: Context value
|
||||
description: Optional description of the context
|
||||
|
||||
Returns:
|
||||
Success status
|
||||
"""
|
||||
logger.info(f"Bridge {self.config.instance_id}: Sharing context '{key}'")
|
||||
|
||||
self.shared_context[key] = {
|
||||
"value": value,
|
||||
"description": description,
|
||||
"shared_by": self.config.source_instance,
|
||||
"shared_with": self.config.target_instance,
|
||||
}
|
||||
|
||||
return True
|
||||
|
||||
@self.tool()
|
||||
async def get_shared_context(key: Optional[str] = None) -> Dict[str, Any]:
|
||||
"""Get shared context from Claude network.
|
||||
|
||||
Args:
|
||||
key: Optional specific key to retrieve
|
||||
|
||||
Returns:
|
||||
Shared context data
|
||||
"""
|
||||
if key:
|
||||
return self.shared_context.get(key, {})
|
||||
return self.shared_context
|
||||
|
||||
@self.tool()
|
||||
async def brainstorm_with_claude(
|
||||
topic: str, num_ideas: int = 5, constraints: Optional[List[str]] = None
|
||||
) -> List[str]:
|
||||
"""Brainstorm ideas with another Claude.
|
||||
|
||||
Args:
|
||||
topic: Topic to brainstorm about
|
||||
num_ideas: Number of ideas to generate
|
||||
constraints: Optional constraints
|
||||
|
||||
Returns:
|
||||
List of brainstormed ideas
|
||||
"""
|
||||
logger.info(f"Bridge {self.config.instance_id}: Brainstorming request")
|
||||
|
||||
brainstorm_prompt = f"""
|
||||
Please brainstorm {num_ideas} ideas about: {topic}
|
||||
|
||||
{"Constraints: " + ", ".join(constraints) if constraints else ""}
|
||||
|
||||
Provide creative and practical ideas.
|
||||
"""
|
||||
|
||||
response = await self._forward_to_claude(brainstorm_prompt)
|
||||
|
||||
# Parse response into list (simplified)
|
||||
ideas = response.split("\n")
|
||||
ideas = [idea.strip() for idea in ideas if idea.strip()]
|
||||
|
||||
return ideas[:num_ideas]
|
||||
|
||||
@self.tool()
|
||||
async def get_claude_status() -> Dict[str, Any]:
|
||||
"""Get status of the connected Claude instance.
|
||||
|
||||
Returns:
|
||||
Status information
|
||||
"""
|
||||
return {
|
||||
"instance_id": self.config.instance_id,
|
||||
"role": self.config.role,
|
||||
"status": "available",
|
||||
"conversation_count": len(self.conversation_history),
|
||||
"shared_context_keys": list(self.shared_context.keys()),
|
||||
}
|
||||
|
||||
def _build_review_prompt(
|
||||
self, code: str, description: str, focus_areas: Optional[List[str]]
|
||||
) -> str:
|
||||
"""Build a code review prompt."""
|
||||
prompt = f"""
|
||||
Please review the following code:
|
||||
|
||||
Description: {description}
|
||||
|
||||
Code:
|
||||
```
|
||||
{code}
|
||||
```
|
||||
"""
|
||||
|
||||
if focus_areas:
|
||||
prompt += f"\n\nPlease focus particularly on: {', '.join(focus_areas)}"
|
||||
|
||||
prompt += """
|
||||
|
||||
Provide constructive feedback on:
|
||||
1. Potential bugs or issues
|
||||
2. Code quality and best practices
|
||||
3. Performance considerations
|
||||
4. Security concerns
|
||||
5. Suggestions for improvement
|
||||
"""
|
||||
|
||||
return prompt
|
||||
|
||||
def _build_delegation_prompt(
|
||||
self, task: str, requirements: List[str], constraints: Optional[List[str]]
|
||||
) -> str:
|
||||
"""Build a task delegation prompt."""
|
||||
prompt = f"""
|
||||
Please complete the following task:
|
||||
|
||||
Task: {task}
|
||||
|
||||
Requirements:
|
||||
{chr(10).join(f"- {req}" for req in requirements)}
|
||||
"""
|
||||
|
||||
if constraints:
|
||||
prompt += f"""
|
||||
|
||||
Constraints:
|
||||
{chr(10).join(f"- {con}" for con in constraints)}
|
||||
"""
|
||||
|
||||
prompt += """
|
||||
|
||||
Provide a complete solution that meets all requirements.
|
||||
"""
|
||||
|
||||
return prompt
|
||||
|
||||
def _build_opinion_prompt(
|
||||
self, question: str, options: Optional[List[str]], criteria: Optional[List[str]]
|
||||
) -> str:
|
||||
"""Build an opinion request prompt."""
|
||||
prompt = f"""
|
||||
I need your opinion on the following:
|
||||
|
||||
Question: {question}
|
||||
"""
|
||||
|
||||
if options:
|
||||
prompt += f"""
|
||||
|
||||
Options to consider:
|
||||
{chr(10).join(f"{i+1}. {opt}" for i, opt in enumerate(options))}
|
||||
"""
|
||||
|
||||
if criteria:
|
||||
prompt += f"""
|
||||
|
||||
Please evaluate based on these criteria:
|
||||
{chr(10).join(f"- {crit}" for crit in criteria)}
|
||||
"""
|
||||
|
||||
prompt += """
|
||||
|
||||
Provide your recommendation with clear reasoning.
|
||||
"""
|
||||
|
||||
return prompt
|
||||
|
||||
async def _forward_to_claude(
|
||||
self, prompt: str, context: Optional[str] = None
|
||||
) -> str:
|
||||
"""Forward a request to the target Claude instance.
|
||||
|
||||
In production, this would make an actual API call to the Claude instance.
|
||||
For now, it returns a simulated response.
|
||||
"""
|
||||
# Add context if provided
|
||||
full_prompt = prompt
|
||||
if context:
|
||||
full_prompt = f"Context: {context}\n\n{prompt}"
|
||||
|
||||
# Log the forwarding
|
||||
logger.info(
|
||||
f"Forwarding from instance {self.config.source_instance} to {self.config.target_instance}"
|
||||
)
|
||||
logger.debug(f"Prompt: {full_prompt[:200]}...")
|
||||
|
||||
# In production, this would:
|
||||
# 1. Connect to the target Claude instance API
|
||||
# 2. Send the prompt
|
||||
# 3. Receive and return the response
|
||||
|
||||
# Simulated response based on role
|
||||
if self.config.role.startswith("critic"):
|
||||
return f"""
|
||||
As {self.config.role}, I've analyzed your request:
|
||||
|
||||
Strengths:
|
||||
- The approach is logical and well-structured
|
||||
- Good attention to requirements
|
||||
|
||||
Areas for improvement:
|
||||
- Consider edge cases more thoroughly
|
||||
- Add more comprehensive error handling
|
||||
- Optimize for performance in high-load scenarios
|
||||
|
||||
Recommendation: Proceed with suggested improvements.
|
||||
"""
|
||||
else:
|
||||
return f"""
|
||||
Response from {self.config.role} (instance {self.config.instance_id}):
|
||||
|
||||
I've processed your request: "{prompt[:100]}..."
|
||||
|
||||
The task has been completed successfully with the following approach:
|
||||
1. Analyzed the requirements
|
||||
2. Implemented the solution
|
||||
3. Validated the results
|
||||
|
||||
The solution meets all specified criteria.
|
||||
"""
|
||||
|
||||
|
||||
async def run_bridge_server(config: BridgeConfig):
|
||||
"""Run the MCP bridge server.
|
||||
|
||||
Args:
|
||||
config: Bridge configuration
|
||||
"""
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
|
||||
logger.info(f"Starting MCP Bridge for Claude instance {config.instance_id}")
|
||||
logger.info(f"Role: {config.role}")
|
||||
logger.info(f"Target port: {config.target_port}")
|
||||
|
||||
# Create and run the bridge
|
||||
bridge = ClaudeBridge(config)
|
||||
|
||||
# Run the stdio server
|
||||
async with stdio_server() as (read_stream, write_stream):
|
||||
await bridge.run(
|
||||
read_stream=read_stream,
|
||||
write_stream=write_stream,
|
||||
initialization_options=InitializationOptions(
|
||||
server_name=bridge.name,
|
||||
server_version="1.0.0",
|
||||
capabilities=bridge.get_capabilities(),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point for the bridge."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="MCP Bridge for Claude-to-Claude communication"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-port",
|
||||
type=int,
|
||||
required=True,
|
||||
help="Port of the target Claude instance",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--instance-id",
|
||||
type=int,
|
||||
required=True,
|
||||
help="ID of the target Claude instance",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--role",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Role of the target instance (primary, critic_1, etc.)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get source/target from environment
|
||||
source_instance = int(os.environ.get("SOURCE_INSTANCE", "0"))
|
||||
target_instance = int(os.environ.get("TARGET_INSTANCE", args.instance_id))
|
||||
|
||||
config = BridgeConfig(
|
||||
target_port=args.target_port,
|
||||
instance_id=args.instance_id,
|
||||
role=args.role,
|
||||
source_instance=source_instance,
|
||||
target_instance=target_instance,
|
||||
)
|
||||
|
||||
# Run the bridge
|
||||
asyncio.run(run_bridge_server(config))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,448 +0,0 @@
|
||||
"""Command-line interface for the Hanzo AI server.
|
||||
|
||||
This module intentionally defers heavy imports (like the server and its
|
||||
dependencies) until after we determine the transport and configure logging.
|
||||
This prevents any stdout/stderr noise from imports that would corrupt the
|
||||
MCP stdio transport used by Claude Desktop and other MCP clients.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import signal
|
||||
import logging
|
||||
import argparse
|
||||
from typing import Any, cast
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Run the CLI for the Hanzo AI server."""
|
||||
# Pre-parse arguments to check transport type early, BEFORE importing server
|
||||
early_parser = argparse.ArgumentParser(add_help=False)
|
||||
early_parser.add_argument("--transport", choices=["stdio", "sse"], default="stdio")
|
||||
early_args, _ = early_parser.parse_known_args()
|
||||
|
||||
# Configure logging VERY early based on transport
|
||||
suppress_stdout = False
|
||||
original_stdout = sys.stdout
|
||||
if early_args.transport == "stdio":
|
||||
# Set environment variable for server to detect stdio mode as early as possible
|
||||
os.environ["HANZO_MCP_TRANSPORT"] = "stdio"
|
||||
# Aggressively quiet common dependency loggers/warnings in stdio mode
|
||||
os.environ.setdefault("PYTHONWARNINGS", "ignore")
|
||||
os.environ.setdefault("LITELLM_LOG", "ERROR")
|
||||
os.environ.setdefault("LITELLM_LOGGING_LEVEL", "ERROR")
|
||||
os.environ.setdefault("FASTMCP_LOG_LEVEL", "ERROR")
|
||||
|
||||
# Suppress FastMCP logging (if available) and all standard logging
|
||||
try:
|
||||
from fastmcp.utilities.logging import configure_logging # type: ignore
|
||||
|
||||
configure_logging(level="ERROR")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.ERROR, # Only show errors
|
||||
handlers=[], # No handlers for stdio to prevent protocol corruption
|
||||
)
|
||||
|
||||
# Redirect stderr to devnull for stdio transport to prevent any output
|
||||
sys.stderr = open(os.devnull, "w")
|
||||
|
||||
# Suppress stdout during potentially noisy imports unless user requested help/version
|
||||
if not any(flag in sys.argv for flag in ("--version", "-h", "--help")):
|
||||
sys.stdout = open(os.devnull, "w")
|
||||
suppress_stdout = True
|
||||
|
||||
# Import the server only AFTER transport/logging have been configured to avoid import-time noise
|
||||
from hanzo_mcp.server import HanzoMCPServer
|
||||
|
||||
# Avoid importing hanzo_mcp package just to get version (it can have side-effects).
|
||||
try:
|
||||
from importlib.metadata import version as _pkg_version # py3.8+
|
||||
_version = _pkg_version("hanzo-mcp")
|
||||
except Exception:
|
||||
_version = "unknown"
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="MCP server implementing Hanzo AI capabilities"
|
||||
)
|
||||
|
||||
parser.add_argument("--version", action="version", version=f"hanzo-mcp {_version}")
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--transport",
|
||||
choices=["stdio", "sse"],
|
||||
default="stdio",
|
||||
help="Transport protocol to use (default: stdio)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--name",
|
||||
default="hanzo-mcp",
|
||||
help="Name of the MCP server (default: hanzo-mcp)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--allow-path",
|
||||
action="append",
|
||||
dest="allowed_paths",
|
||||
help="Add an allowed path (can be specified multiple times)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--project",
|
||||
action="append",
|
||||
dest="project_paths",
|
||||
help="Add a project path for prompt generation (can be specified multiple times)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--agent-model",
|
||||
dest="agent_model",
|
||||
help="Specify the model name in LiteLLM format (e.g., 'openai/gpt-4o', 'anthropic/claude-4-sonnet')",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--agent-max-tokens",
|
||||
dest="agent_max_tokens",
|
||||
type=int,
|
||||
help="Specify the maximum tokens for agent responses",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--agent-api-key",
|
||||
dest="agent_api_key",
|
||||
help="Specify the API key for the LLM provider (for development/testing only)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--agent-base-url",
|
||||
dest="agent_base_url",
|
||||
help="Specify the base URL for the LLM provider API endpoint (e.g., 'http://localhost:1234/v1')",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--agent-max-iterations",
|
||||
dest="agent_max_iterations",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Maximum number of iterations for agent (default: 10)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--agent-max-tool-uses",
|
||||
dest="agent_max_tool_uses",
|
||||
type=int,
|
||||
default=30,
|
||||
help="Maximum number of total tool uses for agent (default: 30)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--enable-agent-tool",
|
||||
dest="enable_agent_tool",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable the agent tool (disabled by default)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--command-timeout",
|
||||
dest="command_timeout",
|
||||
type=float,
|
||||
default=120.0,
|
||||
help="Default timeout for command execution in seconds (default: 120.0)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--disable-write-tools",
|
||||
dest="disable_write_tools",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Disable write tools (edit, write, etc.)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--disable-search-tools",
|
||||
dest="disable_search_tools",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Disable search tools (grep, search_content, etc.)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--host",
|
||||
dest="host",
|
||||
default="127.0.0.1",
|
||||
help="Host for SSE server (default: 127.0.0.1)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--port",
|
||||
dest="port",
|
||||
type=int,
|
||||
default=8888,
|
||||
help="Port for SSE server (default: 8888)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--log-level",
|
||||
dest="log_level",
|
||||
default="INFO",
|
||||
choices=["DEBUG", "INFO", "WARNING", "ERROR"],
|
||||
help="Logging level (default: INFO)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--project-dir",
|
||||
dest="project_dir",
|
||||
help="Single project directory (alias for --project)",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--dev",
|
||||
action="store_true",
|
||||
help="Run in development mode with hot reload",
|
||||
)
|
||||
|
||||
_ = parser.add_argument(
|
||||
"--install",
|
||||
action="store_true",
|
||||
help="Install server configuration in Claude Desktop",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Restore stdout after parsing, before any explicit output or server start
|
||||
if suppress_stdout:
|
||||
try:
|
||||
sys.stdout.close() # Close devnull handle
|
||||
except Exception:
|
||||
pass
|
||||
sys.stdout = original_stdout
|
||||
|
||||
# Cast args attributes to appropriate types to avoid 'Any' warnings
|
||||
name: str = cast(str, args.name)
|
||||
install: bool = cast(bool, args.install)
|
||||
dev: bool = cast(bool, args.dev)
|
||||
transport: str = cast(str, args.transport)
|
||||
agent_model: str | None = cast(str | None, args.agent_model)
|
||||
agent_max_tokens: int | None = cast(int | None, args.agent_max_tokens)
|
||||
agent_api_key: str | None = cast(str | None, args.agent_api_key)
|
||||
agent_base_url: str | None = cast(str | None, args.agent_base_url)
|
||||
agent_max_iterations: int = cast(int, args.agent_max_iterations)
|
||||
agent_max_tool_uses: int = cast(int, args.agent_max_tool_uses)
|
||||
enable_agent_tool: bool = cast(bool, args.enable_agent_tool)
|
||||
command_timeout: float = cast(float, args.command_timeout)
|
||||
disable_write_tools: bool = cast(bool, args.disable_write_tools)
|
||||
disable_search_tools: bool = cast(bool, args.disable_search_tools)
|
||||
host: str = cast(str, args.host)
|
||||
port: int = cast(int, args.port)
|
||||
log_level: str = cast(str, args.log_level)
|
||||
project_dir: str | None = cast(str | None, args.project_dir)
|
||||
allowed_paths: list[str] = (
|
||||
cast(list[str], args.allowed_paths) if args.allowed_paths else []
|
||||
)
|
||||
project_paths: list[str] = (
|
||||
cast(list[str], args.project_paths) if args.project_paths else []
|
||||
)
|
||||
|
||||
# Handle project_dir parameter (add to both allowed_paths and project_paths)
|
||||
if project_dir:
|
||||
if project_dir not in allowed_paths:
|
||||
allowed_paths.append(project_dir)
|
||||
if project_dir not in project_paths:
|
||||
project_paths.append(project_dir)
|
||||
|
||||
if install:
|
||||
install_claude_desktop_config(
|
||||
name, allowed_paths, disable_write_tools, disable_search_tools, host, port
|
||||
)
|
||||
return
|
||||
|
||||
# Get logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Set up signal handler to ensure clean exit
|
||||
def signal_handler(signum, frame):
|
||||
if transport != "stdio":
|
||||
logger.info("\nReceived interrupt signal, shutting down...")
|
||||
sys.exit(0)
|
||||
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
signal.signal(signal.SIGTERM, signal_handler)
|
||||
|
||||
# Configure logging based on transport (stdio already configured early)
|
||||
if transport != "stdio":
|
||||
# For SSE transport, logging is fine
|
||||
log_level_map = {
|
||||
"DEBUG": logging.DEBUG,
|
||||
"INFO": logging.INFO,
|
||||
"WARNING": logging.WARNING,
|
||||
"ERROR": logging.ERROR,
|
||||
}
|
||||
logging.basicConfig(
|
||||
level=log_level_map.get(log_level, logging.INFO),
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
|
||||
# If no allowed paths are specified, use the current directory
|
||||
if not allowed_paths:
|
||||
allowed_paths = [os.getcwd()]
|
||||
|
||||
# Run in dev mode if requested
|
||||
if dev:
|
||||
from hanzo_mcp.dev_server import DevServer
|
||||
|
||||
dev_server = DevServer(
|
||||
name=name,
|
||||
allowed_paths=allowed_paths,
|
||||
project_paths=project_paths,
|
||||
project_dir=project_dir,
|
||||
agent_model=agent_model,
|
||||
agent_max_tokens=agent_max_tokens,
|
||||
agent_api_key=agent_api_key,
|
||||
agent_base_url=agent_base_url,
|
||||
agent_max_iterations=agent_max_iterations,
|
||||
agent_max_tool_uses=agent_max_tool_uses,
|
||||
enable_agent_tool=enable_agent_tool,
|
||||
command_timeout=command_timeout,
|
||||
disable_write_tools=disable_write_tools,
|
||||
disable_search_tools=disable_search_tools,
|
||||
host=host,
|
||||
port=port,
|
||||
)
|
||||
dev_server.run(transport=transport)
|
||||
return
|
||||
|
||||
# Run the server
|
||||
server = HanzoMCPServer(
|
||||
name=name,
|
||||
allowed_paths=allowed_paths,
|
||||
project_paths=project_paths,
|
||||
project_dir=project_dir,
|
||||
agent_model=agent_model,
|
||||
agent_max_tokens=agent_max_tokens,
|
||||
agent_api_key=agent_api_key,
|
||||
agent_base_url=agent_base_url,
|
||||
agent_max_iterations=agent_max_iterations,
|
||||
agent_max_tool_uses=agent_max_tool_uses,
|
||||
enable_agent_tool=enable_agent_tool,
|
||||
command_timeout=command_timeout,
|
||||
disable_write_tools=disable_write_tools,
|
||||
disable_search_tools=disable_search_tools,
|
||||
host=host,
|
||||
port=port,
|
||||
)
|
||||
|
||||
try:
|
||||
# Transport will be automatically cast to Literal['stdio', 'sse'] by the server
|
||||
server.run(transport=transport)
|
||||
except KeyboardInterrupt:
|
||||
if transport != "stdio":
|
||||
logger.info("\nShutting down...")
|
||||
sys.exit(0)
|
||||
except Exception as e:
|
||||
logger.error(f"Server error: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def install_claude_desktop_config(
|
||||
name: str = "hanzo-mcp",
|
||||
allowed_paths: list[str] | None = None,
|
||||
disable_write_tools: bool = False,
|
||||
disable_search_tools: bool = False,
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 8888,
|
||||
) -> None:
|
||||
"""Install the server configuration in Claude Desktop.
|
||||
|
||||
Args:
|
||||
name: The name to use for the server in the config
|
||||
allowed_paths: Optional list of paths to allow
|
||||
disable_write_tools: Whether to disable write tools
|
||||
disable_search_tools: Whether to disable search tools
|
||||
host: Host for SSE server
|
||||
port: Port for SSE server
|
||||
"""
|
||||
# Find the Claude Desktop config directory
|
||||
home: Path = Path.home()
|
||||
|
||||
if sys.platform == "darwin": # macOS
|
||||
config_dir: Path = home / "Library" / "Application Support" / "Claude"
|
||||
elif sys.platform == "win32": # Windows
|
||||
config_dir = Path(os.environ.get("APPDATA", "")) / "Claude"
|
||||
else: # Linux and others
|
||||
config_dir = home / ".config" / "claude"
|
||||
|
||||
config_file: Path = config_dir / "claude_desktop_config.json"
|
||||
|
||||
# Create directory if it doesn't exist
|
||||
config_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Get current script path
|
||||
script_path: Path = Path(sys.executable)
|
||||
|
||||
# Create args array
|
||||
args: list[str] = ["-m", "hanzo_mcp.cli"]
|
||||
|
||||
# Add allowed paths if specified
|
||||
if allowed_paths:
|
||||
for path in allowed_paths:
|
||||
args.extend(["--allow-path", path])
|
||||
else:
|
||||
# Allow home directory by default
|
||||
args.extend(["--allow-path", str(home)])
|
||||
|
||||
# Add tool disable flags if specified
|
||||
if disable_write_tools:
|
||||
args.append("--disable-write-tools")
|
||||
|
||||
if disable_search_tools:
|
||||
args.append("--disable-search-tools")
|
||||
|
||||
# Create config object
|
||||
config: dict[str, Any] = {
|
||||
"mcpServers": {name: {"command": script_path.as_posix(), "args": args}}
|
||||
}
|
||||
|
||||
# Check if the file already exists
|
||||
if config_file.exists():
|
||||
try:
|
||||
with open(config_file, "r") as f:
|
||||
existing_config: dict[str, Any] = json.load(f)
|
||||
|
||||
# Update the existing config
|
||||
if "mcpServers" not in existing_config:
|
||||
existing_config["mcpServers"] = {}
|
||||
|
||||
existing_config["mcpServers"][name] = config["mcpServers"][name]
|
||||
config = existing_config
|
||||
except Exception as e:
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.error(f"Error reading existing config: {e}")
|
||||
logger.info("Creating new config file.")
|
||||
|
||||
# Write the config file
|
||||
with open(config_file, mode="w") as f:
|
||||
json.dump(config, f, indent=2)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(f"Successfully installed {name} in Claude Desktop configuration.")
|
||||
logger.info(f"Config file: {config_file}")
|
||||
|
||||
if allowed_paths:
|
||||
logger.info("\nAllowed paths:")
|
||||
for path in allowed_paths:
|
||||
logger.info(f"- {path}")
|
||||
else:
|
||||
logger.info(f"\nDefault allowed path: {home}")
|
||||
|
||||
logger.info("\nYou can modify allowed paths in the config file directly.")
|
||||
logger.info("Restart Claude Desktop for changes to take effect.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,463 +0,0 @@
|
||||
"""Enhanced command-line interface for the Hanzo AI server with full tool configuration."""
|
||||
|
||||
import os
|
||||
import logging
|
||||
import argparse
|
||||
from typing import Any, Dict
|
||||
|
||||
from hanzo_mcp.config import (
|
||||
TOOL_REGISTRY,
|
||||
HanzoMCPSettings,
|
||||
load_settings,
|
||||
save_settings,
|
||||
)
|
||||
from hanzo_mcp.server import HanzoMCPServer
|
||||
|
||||
|
||||
def create_parser() -> argparse.ArgumentParser:
|
||||
"""Create the argument parser with all tool configuration options."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Hanzo AI server with comprehensive tool configuration",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Tool Configuration:
|
||||
Each tool can be individually enabled/disabled using CLI flags.
|
||||
Use --list-tools to see all available tools and their current status.
|
||||
|
||||
Configuration Files:
|
||||
Global: ~/.config/hanzo/mcp-settings.json
|
||||
Project: ./.hanzo-mcp.json or ./.hanzo/mcp-settings.json
|
||||
|
||||
Examples:
|
||||
# Start server with only read tools
|
||||
hanzo-mcp --disable-write --disable-edit --disable-multi-edit
|
||||
|
||||
# Enable agent tool with custom model
|
||||
hanzo-mcp --enable-dispatch-agent --agent-model anthropic/claude-3-sonnet
|
||||
|
||||
# Save current configuration
|
||||
hanzo-mcp --save-config
|
||||
""",
|
||||
)
|
||||
|
||||
# Basic server options
|
||||
server_group = parser.add_argument_group("Server Configuration")
|
||||
server_group.add_argument(
|
||||
"--name",
|
||||
default="hanzo-mcp",
|
||||
help="Name of the MCP server (default: hanzo-mcp)",
|
||||
)
|
||||
server_group.add_argument(
|
||||
"--transport",
|
||||
choices=["stdio", "sse"],
|
||||
default="stdio",
|
||||
help="Transport protocol to use (default: stdio)",
|
||||
)
|
||||
server_group.add_argument(
|
||||
"--host",
|
||||
default="127.0.0.1",
|
||||
help="Host for SSE server (default: 127.0.0.1)",
|
||||
)
|
||||
server_group.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=8888,
|
||||
help="Port for SSE server (default: 8888)",
|
||||
)
|
||||
server_group.add_argument(
|
||||
"--log-level",
|
||||
choices=["DEBUG", "INFO", "WARNING", "ERROR"],
|
||||
default="INFO",
|
||||
help="Logging level (default: INFO)",
|
||||
)
|
||||
server_group.add_argument(
|
||||
"--command-timeout",
|
||||
type=float,
|
||||
default=120.0,
|
||||
help="Default timeout for command execution in seconds (default: 120.0)",
|
||||
)
|
||||
|
||||
# Path configuration
|
||||
path_group = parser.add_argument_group("Path Configuration")
|
||||
path_group.add_argument(
|
||||
"--allow-path",
|
||||
action="append",
|
||||
dest="allowed_paths",
|
||||
help="Add an allowed path (can be specified multiple times)",
|
||||
)
|
||||
path_group.add_argument(
|
||||
"--project",
|
||||
action="append",
|
||||
dest="project_paths",
|
||||
help="Add a project path (can be specified multiple times)",
|
||||
)
|
||||
path_group.add_argument(
|
||||
"--project-dir",
|
||||
help="Single project directory (added to both allowed and project paths)",
|
||||
)
|
||||
|
||||
# Individual tool configuration
|
||||
tool_group = parser.add_argument_group("Individual Tool Configuration")
|
||||
|
||||
# Add CLI flags for each tool
|
||||
for tool_name, tool_config in TOOL_REGISTRY.items():
|
||||
flag_name = tool_config.cli_flag.lstrip("-")
|
||||
help_text = f"{tool_config.description}"
|
||||
|
||||
if tool_config.enabled:
|
||||
# Tool is enabled by default, add disable flag
|
||||
tool_group.add_argument(
|
||||
f"--disable-{tool_name.replace('_', '-')}",
|
||||
action="store_true",
|
||||
help=f"Disable {tool_name} tool: {help_text}",
|
||||
)
|
||||
else:
|
||||
# Tool is disabled by default, add enable flag
|
||||
tool_group.add_argument(
|
||||
f"--enable-{tool_name.replace('_', '-')}",
|
||||
action="store_true",
|
||||
help=f"Enable {tool_name} tool: {help_text}",
|
||||
)
|
||||
|
||||
# Category-level tool configuration (for backward compatibility)
|
||||
category_group = parser.add_argument_group("Category-level Tool Configuration")
|
||||
category_group.add_argument(
|
||||
"--disable-write-tools",
|
||||
action="store_true",
|
||||
help="Disable all write tools (write, edit, multi_edit, content_replace)",
|
||||
)
|
||||
category_group.add_argument(
|
||||
"--disable-search-tools",
|
||||
action="store_true",
|
||||
help="Disable all search tools (grep, grep_ast)",
|
||||
)
|
||||
category_group.add_argument(
|
||||
"--disable-filesystem-tools",
|
||||
action="store_true",
|
||||
help="Disable all filesystem tools",
|
||||
)
|
||||
category_group.add_argument(
|
||||
"--disable-jupyter-tools",
|
||||
action="store_true",
|
||||
help="Disable all Jupyter notebook tools",
|
||||
)
|
||||
category_group.add_argument(
|
||||
"--disable-shell-tools",
|
||||
action="store_true",
|
||||
help="Disable shell command execution tools",
|
||||
)
|
||||
category_group.add_argument(
|
||||
"--disable-todo-tools",
|
||||
action="store_true",
|
||||
help="Disable todo management tools",
|
||||
)
|
||||
|
||||
# Agent configuration
|
||||
agent_group = parser.add_argument_group("Agent Tool Configuration")
|
||||
agent_group.add_argument(
|
||||
"--agent-model",
|
||||
help="Model name in LiteLLM format (e.g., 'openai/gpt-4o', 'anthropic/claude-3-sonnet')",
|
||||
)
|
||||
agent_group.add_argument(
|
||||
"--agent-max-tokens",
|
||||
type=int,
|
||||
help="Maximum tokens for agent responses",
|
||||
)
|
||||
agent_group.add_argument(
|
||||
"--agent-api-key",
|
||||
help="API key for the LLM provider",
|
||||
)
|
||||
agent_group.add_argument(
|
||||
"--agent-base-url",
|
||||
help="Base URL for the LLM provider API endpoint",
|
||||
)
|
||||
agent_group.add_argument(
|
||||
"--agent-max-iterations",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Maximum iterations for agent (default: 10)",
|
||||
)
|
||||
agent_group.add_argument(
|
||||
"--agent-max-tool-uses",
|
||||
type=int,
|
||||
default=30,
|
||||
help="Maximum tool uses for agent (default: 30)",
|
||||
)
|
||||
|
||||
# Vector store configuration
|
||||
vector_group = parser.add_argument_group("Vector Store Configuration")
|
||||
vector_group.add_argument(
|
||||
"--enable-vector-store",
|
||||
action="store_true",
|
||||
help="Enable local vector store (Infinity database)",
|
||||
)
|
||||
vector_group.add_argument(
|
||||
"--vector-store-path",
|
||||
help="Path for vector store data (default: ~/.config/hanzo/vector-store)",
|
||||
)
|
||||
vector_group.add_argument(
|
||||
"--embedding-model",
|
||||
default="text-embedding-3-small",
|
||||
help="Embedding model for vector store (default: text-embedding-3-small)",
|
||||
)
|
||||
|
||||
# Configuration management
|
||||
config_group = parser.add_argument_group("Configuration Management")
|
||||
config_group.add_argument(
|
||||
"--config-file",
|
||||
help="Load configuration from specific file",
|
||||
)
|
||||
config_group.add_argument(
|
||||
"--save-config",
|
||||
action="store_true",
|
||||
help="Save current configuration to global config file",
|
||||
)
|
||||
config_group.add_argument(
|
||||
"--save-project-config",
|
||||
action="store_true",
|
||||
help="Save current configuration to project config file",
|
||||
)
|
||||
config_group.add_argument(
|
||||
"--list-tools",
|
||||
action="store_true",
|
||||
help="List all available tools and their status",
|
||||
)
|
||||
|
||||
# Installation
|
||||
install_group = parser.add_argument_group("Installation")
|
||||
install_group.add_argument(
|
||||
"--install",
|
||||
action="store_true",
|
||||
help="Install server configuration in Claude Desktop",
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def apply_cli_overrides(args: argparse.Namespace) -> Dict[str, Any]:
|
||||
"""Convert CLI arguments to configuration overrides."""
|
||||
overrides = {}
|
||||
|
||||
# Server configuration
|
||||
server_config = {}
|
||||
if hasattr(args, "name") and args.name != "hanzo-mcp":
|
||||
server_config["name"] = args.name
|
||||
if hasattr(args, "host") and args.host != "127.0.0.1":
|
||||
server_config["host"] = args.host
|
||||
if hasattr(args, "port") and args.port != 8888:
|
||||
server_config["port"] = args.port
|
||||
if hasattr(args, "transport") and args.transport != "stdio":
|
||||
server_config["transport"] = args.transport
|
||||
if hasattr(args, "log_level") and args.log_level != "INFO":
|
||||
server_config["log_level"] = args.log_level
|
||||
if hasattr(args, "command_timeout") and args.command_timeout != 120.0:
|
||||
server_config["command_timeout"] = args.command_timeout
|
||||
|
||||
if server_config:
|
||||
overrides["server"] = server_config
|
||||
|
||||
# Path configuration
|
||||
if hasattr(args, "allowed_paths") and args.allowed_paths:
|
||||
overrides["allowed_paths"] = args.allowed_paths
|
||||
if hasattr(args, "project_paths") and args.project_paths:
|
||||
overrides["project_paths"] = args.project_paths
|
||||
if hasattr(args, "project_dir") and args.project_dir:
|
||||
overrides["project_dir"] = args.project_dir
|
||||
|
||||
# Tool configuration
|
||||
enabled_tools = {}
|
||||
|
||||
# Handle individual tool flags
|
||||
for tool_name, tool_config in TOOL_REGISTRY.items():
|
||||
flag_name = tool_name.replace("_", "-")
|
||||
|
||||
if tool_config.enabled:
|
||||
# Check for disable flag
|
||||
disable_flag = f"disable_{tool_name}"
|
||||
if hasattr(args, disable_flag) and getattr(args, disable_flag):
|
||||
enabled_tools[tool_name] = False
|
||||
else:
|
||||
# Check for enable flag
|
||||
enable_flag = f"enable_{tool_name}"
|
||||
if hasattr(args, enable_flag) and getattr(args, enable_flag):
|
||||
enabled_tools[tool_name] = True
|
||||
|
||||
# Handle category-level disables
|
||||
if hasattr(args, "disable_write_tools") and args.disable_write_tools:
|
||||
for tool_name in ["write", "edit", "multi_edit", "content_replace"]:
|
||||
enabled_tools[tool_name] = False
|
||||
|
||||
if hasattr(args, "disable_search_tools") and args.disable_search_tools:
|
||||
for tool_name in ["grep", "grep_ast"]:
|
||||
enabled_tools[tool_name] = False
|
||||
|
||||
if hasattr(args, "disable_filesystem_tools") and args.disable_filesystem_tools:
|
||||
filesystem_tools = [
|
||||
"read",
|
||||
"write",
|
||||
"edit",
|
||||
"multi_edit",
|
||||
"directory_tree",
|
||||
"grep",
|
||||
"grep_ast",
|
||||
"content_replace",
|
||||
]
|
||||
for tool_name in filesystem_tools:
|
||||
enabled_tools[tool_name] = False
|
||||
|
||||
if hasattr(args, "disable_jupyter_tools") and args.disable_jupyter_tools:
|
||||
for tool_name in ["notebook_read", "notebook_edit"]:
|
||||
enabled_tools[tool_name] = False
|
||||
|
||||
if hasattr(args, "disable_shell_tools") and args.disable_shell_tools:
|
||||
enabled_tools["run_command"] = False
|
||||
|
||||
if hasattr(args, "disable_todo_tools") and args.disable_todo_tools:
|
||||
for tool_name in ["todo_read", "todo_write"]:
|
||||
enabled_tools[tool_name] = False
|
||||
|
||||
if enabled_tools:
|
||||
overrides["enabled_tools"] = enabled_tools
|
||||
|
||||
# Agent configuration
|
||||
agent_config = {}
|
||||
if hasattr(args, "agent_model") and args.agent_model:
|
||||
agent_config["model"] = args.agent_model
|
||||
agent_config["enabled"] = True
|
||||
if hasattr(args, "agent_api_key") and args.agent_api_key:
|
||||
agent_config["api_key"] = args.agent_api_key
|
||||
if hasattr(args, "agent_base_url") and args.agent_base_url:
|
||||
agent_config["base_url"] = args.agent_base_url
|
||||
if hasattr(args, "agent_max_tokens") and args.agent_max_tokens:
|
||||
agent_config["max_tokens"] = args.agent_max_tokens
|
||||
if hasattr(args, "agent_max_iterations") and args.agent_max_iterations != 10:
|
||||
agent_config["max_iterations"] = args.agent_max_iterations
|
||||
if hasattr(args, "agent_max_tool_uses") and args.agent_max_tool_uses != 30:
|
||||
agent_config["max_tool_uses"] = args.agent_max_tool_uses
|
||||
|
||||
if agent_config:
|
||||
overrides["agent"] = agent_config
|
||||
|
||||
# Vector store configuration
|
||||
vector_config = {}
|
||||
if hasattr(args, "enable_vector_store") and args.enable_vector_store:
|
||||
vector_config["enabled"] = True
|
||||
if hasattr(args, "vector_store_path") and args.vector_store_path:
|
||||
vector_config["data_path"] = args.vector_store_path
|
||||
if (
|
||||
hasattr(args, "embedding_model")
|
||||
and args.embedding_model != "text-embedding-3-small"
|
||||
):
|
||||
vector_config["embedding_model"] = args.embedding_model
|
||||
|
||||
if vector_config:
|
||||
overrides["vector_store"] = vector_config
|
||||
|
||||
return overrides
|
||||
|
||||
|
||||
def list_tools(settings: HanzoMCPSettings) -> None:
|
||||
"""List all tools and their current status."""
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info("Hanzo AI Tools Status:")
|
||||
logger.info("=" * 50)
|
||||
|
||||
categories = {}
|
||||
for tool_name, tool_config in TOOL_REGISTRY.items():
|
||||
category = tool_config.category.value
|
||||
if category not in categories:
|
||||
categories[category] = []
|
||||
|
||||
enabled = settings.is_tool_enabled(tool_name)
|
||||
status = "✅ ENABLED " if enabled else "❌ DISABLED"
|
||||
categories[category].append((tool_name, status, tool_config.description))
|
||||
|
||||
for category, tools in categories.items():
|
||||
logger.info(f"\n{category.upper()} TOOLS:")
|
||||
logger.info("-" * 30)
|
||||
for tool_name, status, description in tools:
|
||||
logger.info(f" {status} {tool_name:<15} - {description}")
|
||||
|
||||
logger.info(f"\nTotal: {len(TOOL_REGISTRY)} tools")
|
||||
enabled_count = len(settings.get_enabled_tools())
|
||||
logger.info(
|
||||
f"Enabled: {enabled_count}, Disabled: {len(TOOL_REGISTRY) - enabled_count}"
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Run the enhanced CLI for the Hanzo AI server."""
|
||||
parser = create_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
# Handle list tools command
|
||||
if hasattr(args, "list_tools") and args.list_tools:
|
||||
settings = load_settings()
|
||||
list_tools(settings)
|
||||
return
|
||||
|
||||
# Load configuration with CLI overrides
|
||||
config_overrides = apply_cli_overrides(args)
|
||||
project_dir = getattr(args, "project_dir", None)
|
||||
settings = load_settings(project_dir=project_dir, config_overrides=config_overrides)
|
||||
|
||||
# Handle configuration saving
|
||||
logger = logging.getLogger(__name__)
|
||||
if hasattr(args, "save_config") and args.save_config:
|
||||
saved_path = save_settings(settings, global_config=True)
|
||||
logger.info(f"Configuration saved to: {saved_path}")
|
||||
return
|
||||
|
||||
if hasattr(args, "save_project_config") and args.save_project_config:
|
||||
saved_path = save_settings(settings, global_config=False)
|
||||
logger.info(f"Project configuration saved to: {saved_path}")
|
||||
return
|
||||
|
||||
# Handle installation
|
||||
if hasattr(args, "install") and args.install:
|
||||
from hanzo_mcp.cli import install_claude_desktop_config
|
||||
|
||||
install_claude_desktop_config(
|
||||
settings.server.name,
|
||||
settings.allowed_paths,
|
||||
)
|
||||
return
|
||||
|
||||
# Set up allowed paths
|
||||
allowed_paths = settings.allowed_paths[:]
|
||||
if settings.project_dir and settings.project_dir not in allowed_paths:
|
||||
allowed_paths.append(settings.project_dir)
|
||||
|
||||
if not allowed_paths:
|
||||
allowed_paths = [os.getcwd()]
|
||||
|
||||
# Create and run server
|
||||
server = HanzoMCPServer(
|
||||
name=settings.server.name,
|
||||
allowed_paths=allowed_paths,
|
||||
project_dir=settings.project_dir,
|
||||
agent_model=settings.agent.model,
|
||||
agent_max_tokens=settings.agent.max_tokens,
|
||||
agent_api_key=settings.agent.api_key,
|
||||
agent_base_url=settings.agent.base_url,
|
||||
agent_max_iterations=settings.agent.max_iterations,
|
||||
agent_max_tool_uses=settings.agent.max_tool_uses,
|
||||
enable_agent_tool=settings.agent.enabled
|
||||
or settings.is_tool_enabled("dispatch_agent"),
|
||||
disable_write_tools=not any(
|
||||
settings.is_tool_enabled(t)
|
||||
for t in ["write", "edit", "multi_edit", "content_replace"]
|
||||
),
|
||||
disable_search_tools=not any(
|
||||
settings.is_tool_enabled(t) for t in ["grep", "grep_ast"]
|
||||
),
|
||||
host=settings.server.host,
|
||||
port=settings.server.port,
|
||||
enabled_tools=settings.enabled_tools, # Pass individual tool configuration
|
||||
)
|
||||
|
||||
server.run(transport=settings.server.transport)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,101 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""CLI for managing Hanzo MCP plugins."""
|
||||
|
||||
import sys
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
from hanzo_mcp.tools.common.plugin_loader import (
|
||||
create_plugin_template,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main CLI entry point."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Hanzo MCP Plugin Manager",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
# Create a new plugin template
|
||||
hanzo-plugin create mytool
|
||||
|
||||
# List installed plugins
|
||||
hanzo-plugin list
|
||||
|
||||
# Create plugin in specific directory
|
||||
hanzo-plugin create mytool --output /path/to/plugins
|
||||
""",
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(dest="command", help="Command to run")
|
||||
|
||||
# Create command
|
||||
create_parser = subparsers.add_parser("create", help="Create a new plugin template")
|
||||
create_parser.add_argument("name", help="Name of the tool (e.g., 'mytool')")
|
||||
create_parser.add_argument(
|
||||
"--output",
|
||||
"-o",
|
||||
type=Path,
|
||||
default=Path.home() / ".hanzo" / "plugins",
|
||||
help="Output directory for the plugin (default: ~/.hanzo/plugins)",
|
||||
)
|
||||
|
||||
# List command
|
||||
list_parser = subparsers.add_parser("list", help="List installed plugins")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command == "create":
|
||||
# Create plugin template
|
||||
output_dir = args.output / args.name
|
||||
try:
|
||||
create_plugin_template(output_dir, args.name)
|
||||
print(f"\n✅ Plugin template created successfully!")
|
||||
print(f"\nTo use your plugin:")
|
||||
print(
|
||||
f"1. Edit the tool implementation in {output_dir / f'{args.name}_tool.py'}"
|
||||
)
|
||||
print(f"2. Restart Hanzo MCP to load the plugin")
|
||||
print(f"3. Add '{args.name}' to your mode's tool list")
|
||||
except Exception as e:
|
||||
print(f"❌ Error creating plugin: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
elif args.command == "list":
|
||||
# List installed plugins
|
||||
try:
|
||||
from hanzo_mcp.tools.common.plugin_loader import load_user_plugins
|
||||
|
||||
plugins = load_user_plugins()
|
||||
|
||||
if not plugins:
|
||||
print("No plugins installed.")
|
||||
print("\nPlugin directories:")
|
||||
print(" ~/.hanzo/plugins/")
|
||||
print(" ./.hanzo/plugins/")
|
||||
print(" $HANZO_PLUGIN_PATH")
|
||||
else:
|
||||
print(f"Installed plugins ({len(plugins)}):")
|
||||
for name, plugin in plugins.items():
|
||||
print(f"\n {name}:")
|
||||
print(f" Source: {plugin.source_path}")
|
||||
if plugin.metadata:
|
||||
print(
|
||||
f" Version: {plugin.metadata.get('version', 'unknown')}"
|
||||
)
|
||||
print(f" Author: {plugin.metadata.get('author', 'unknown')}")
|
||||
print(
|
||||
f" Description: {plugin.metadata.get('description', '')}"
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"❌ Error listing plugins: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
else:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,192 +0,0 @@
|
||||
"""Compute node detection and management for distributed processing."""
|
||||
|
||||
import os
|
||||
import platform
|
||||
import subprocess
|
||||
from typing import Any, Dict, List
|
||||
|
||||
|
||||
class ComputeNodeDetector:
|
||||
"""Detect available compute nodes (GPUs, WebGPU, CPUs) for distributed work."""
|
||||
|
||||
@staticmethod
|
||||
def detect_local_gpus() -> List[Dict[str, Any]]:
|
||||
"""Detect local GPU devices."""
|
||||
gpus = []
|
||||
|
||||
# Try NVIDIA GPUs
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"nvidia-smi",
|
||||
"--query-gpu=name,memory.total",
|
||||
"--format=csv,noheader",
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=2,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
for line in result.stdout.strip().split("\n"):
|
||||
if line:
|
||||
name, memory = line.split(", ")
|
||||
gpus.append(
|
||||
{
|
||||
"type": "cuda",
|
||||
"name": name,
|
||||
"memory": memory,
|
||||
"id": f"cuda:{len(gpus)}",
|
||||
}
|
||||
)
|
||||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
|
||||
# Try Metal GPUs (macOS)
|
||||
if platform.system() == "Darwin":
|
||||
try:
|
||||
# Check for Metal support
|
||||
result = subprocess.run(
|
||||
["system_profiler", "SPDisplaysDataType"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=2,
|
||||
)
|
||||
if result.returncode == 0 and "Metal" in result.stdout:
|
||||
# Parse GPU info from system_profiler
|
||||
lines = result.stdout.split("\n")
|
||||
for i, line in enumerate(lines):
|
||||
if "Chipset Model:" in line:
|
||||
gpu_name = line.split(":")[1].strip()
|
||||
gpus.append(
|
||||
{
|
||||
"type": "metal",
|
||||
"name": gpu_name,
|
||||
"memory": "Shared",
|
||||
"id": f"metal:{len(gpus)}",
|
||||
}
|
||||
)
|
||||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
|
||||
return gpus
|
||||
|
||||
@staticmethod
|
||||
def detect_webgpu_nodes() -> List[Dict[str, Any]]:
|
||||
"""Detect connected WebGPU nodes (from browsers)."""
|
||||
webgpu_nodes = []
|
||||
|
||||
# Check for WebGPU connections (would need actual WebSocket/server to track)
|
||||
# For now, check if a WebGPU server is running
|
||||
webgpu_port = os.environ.get("HANZO_WEBGPU_PORT", "8765")
|
||||
try:
|
||||
import socket
|
||||
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
result = sock.connect_ex(("localhost", int(webgpu_port)))
|
||||
sock.close()
|
||||
if result == 0:
|
||||
webgpu_nodes.append(
|
||||
{
|
||||
"type": "webgpu",
|
||||
"name": "Chrome WebGPU",
|
||||
"memory": "Browser",
|
||||
"id": "webgpu:0",
|
||||
}
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return webgpu_nodes
|
||||
|
||||
@staticmethod
|
||||
def detect_cpu_nodes() -> List[Dict[str, Any]]:
|
||||
"""Detect CPU compute nodes."""
|
||||
import multiprocessing
|
||||
|
||||
return [
|
||||
{
|
||||
"type": "cpu",
|
||||
"name": f"{platform.processor() or 'CPU'}",
|
||||
"cores": multiprocessing.cpu_count(),
|
||||
"id": "cpu:0",
|
||||
}
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def get_all_nodes(cls) -> List[Dict[str, Any]]:
|
||||
"""Get all available compute nodes."""
|
||||
nodes = []
|
||||
|
||||
# Detect GPUs
|
||||
gpus = cls.detect_local_gpus()
|
||||
nodes.extend(gpus)
|
||||
|
||||
# Detect WebGPU connections
|
||||
webgpu = cls.detect_webgpu_nodes()
|
||||
nodes.extend(webgpu)
|
||||
|
||||
# If no GPUs/WebGPU, add CPU as compute node
|
||||
if not nodes:
|
||||
nodes.extend(cls.detect_cpu_nodes())
|
||||
|
||||
return nodes
|
||||
|
||||
@classmethod
|
||||
def get_node_count(cls) -> int:
|
||||
"""Get total number of available compute nodes."""
|
||||
return len(cls.get_all_nodes())
|
||||
|
||||
@classmethod
|
||||
def get_node_summary(cls) -> str:
|
||||
"""Get a summary string of available nodes."""
|
||||
nodes = cls.get_all_nodes()
|
||||
if not nodes:
|
||||
return "No compute nodes available"
|
||||
|
||||
count = len(nodes)
|
||||
node_word = "node" if count == 1 else "nodes"
|
||||
|
||||
# Group by type
|
||||
types = {}
|
||||
for node in nodes:
|
||||
node_type = node["type"]
|
||||
if node_type not in types:
|
||||
types[node_type] = 0
|
||||
types[node_type] += 1
|
||||
|
||||
# Build summary
|
||||
parts = []
|
||||
for node_type, type_count in types.items():
|
||||
if node_type == "cuda":
|
||||
parts.append(f"{type_count} CUDA GPU{'s' if type_count > 1 else ''}")
|
||||
elif node_type == "metal":
|
||||
parts.append(f"{type_count} Metal GPU{'s' if type_count > 1 else ''}")
|
||||
elif node_type == "webgpu":
|
||||
parts.append(f"{type_count} WebGPU")
|
||||
elif node_type == "cpu":
|
||||
parts.append(f"{type_count} CPU")
|
||||
|
||||
type_str = ", ".join(parts)
|
||||
return f"{count} {node_word} available ({type_str})"
|
||||
|
||||
|
||||
def print_node_status():
|
||||
"""Print current node status."""
|
||||
detector = ComputeNodeDetector()
|
||||
nodes = detector.get_all_nodes()
|
||||
|
||||
print(f"\n🖥️ Compute Nodes: {len(nodes)}")
|
||||
for node in nodes:
|
||||
if node["type"] in ["cuda", "metal"]:
|
||||
print(f" • {node['id']}: {node['name']} ({node['memory']})")
|
||||
elif node["type"] == "webgpu":
|
||||
print(f" • {node['id']}: {node['name']}")
|
||||
elif node["type"] == "cpu":
|
||||
print(f" • {node['id']}: {node['name']} ({node['cores']} cores)")
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Test the detector
|
||||
print_node_status()
|
||||
print(ComputeNodeDetector.get_node_summary())
|
||||
@@ -1,19 +0,0 @@
|
||||
"""Configuration management for Hanzo AI.
|
||||
|
||||
This module provides a comprehensive configuration system that supports:
|
||||
- CLI arguments for individual tool control
|
||||
- Configuration files for persistent settings
|
||||
- Environment variable overrides
|
||||
- Per-project settings
|
||||
"""
|
||||
|
||||
from .settings import HanzoMCPSettings, load_settings, save_settings
|
||||
from .tool_config import ToolConfig, ToolCategory
|
||||
|
||||
__all__ = [
|
||||
"HanzoMCPSettings",
|
||||
"ToolConfig",
|
||||
"ToolCategory",
|
||||
"load_settings",
|
||||
"save_settings",
|
||||
]
|
||||
@@ -1,551 +0,0 @@
|
||||
"""Settings management for Hanzo AI.
|
||||
|
||||
Handles loading and saving configuration from multiple sources:
|
||||
1. Default settings
|
||||
2. Global config file (~/.config/hanzo/mcp-settings.json)
|
||||
3. Project-specific config file
|
||||
4. Environment variables
|
||||
5. CLI arguments (highest priority)
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional
|
||||
from pathlib import Path
|
||||
from dataclasses import field, asdict, dataclass
|
||||
|
||||
from .tool_config import TOOL_REGISTRY
|
||||
|
||||
|
||||
@dataclass
|
||||
class MCPServerConfig:
|
||||
"""Configuration for an external MCP server."""
|
||||
|
||||
name: str
|
||||
command: str
|
||||
args: List[str] = field(default_factory=list)
|
||||
enabled: bool = True
|
||||
trusted: bool = False # Whether this server is trusted
|
||||
description: str = ""
|
||||
capabilities: List[str] = field(default_factory=list) # tools, prompts, resources
|
||||
|
||||
|
||||
@dataclass
|
||||
class VectorStoreConfig:
|
||||
"""Configuration for the vector store."""
|
||||
|
||||
enabled: bool = False
|
||||
provider: str = "infinity" # infinity, chroma, etc.
|
||||
data_path: Optional[str] = None # Will default to ~/.config/hanzo/vector-store
|
||||
embedding_model: str = "text-embedding-3-small"
|
||||
chunk_size: int = 1000
|
||||
chunk_overlap: int = 200
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProjectConfig:
|
||||
"""Configuration specific to a project."""
|
||||
|
||||
name: str
|
||||
root_path: str
|
||||
rules: List[str] = field(default_factory=list)
|
||||
workflows: Dict[str, Any] = field(default_factory=dict)
|
||||
tasks: List[Dict[str, Any]] = field(default_factory=list)
|
||||
enabled_tools: Dict[str, bool] = field(default_factory=dict)
|
||||
disabled_tools: List[str] = field(default_factory=list)
|
||||
mcp_servers: List[str] = field(
|
||||
default_factory=list
|
||||
) # Names of enabled MCP servers for this project
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentConfig:
|
||||
"""Configuration for agent tools."""
|
||||
|
||||
enabled: bool = False
|
||||
model: Optional[str] = None
|
||||
api_key: Optional[str] = None
|
||||
hanzo_api_key: Optional[str] = None # HANZO_API_KEY support
|
||||
base_url: Optional[str] = None
|
||||
max_tokens: Optional[int] = None
|
||||
max_iterations: int = 10
|
||||
max_tool_uses: int = 30
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServerConfig:
|
||||
"""Configuration for the MCP server."""
|
||||
|
||||
name: str = "hanzo-mcp"
|
||||
host: str = "127.0.0.1"
|
||||
port: int = 8888
|
||||
transport: str = "stdio" # stdio or sse
|
||||
log_level: str = "INFO"
|
||||
command_timeout: float = 120.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class HanzoMCPSettings:
|
||||
"""Complete configuration for Hanzo AI."""
|
||||
|
||||
# Server settings
|
||||
server: ServerConfig = field(default_factory=ServerConfig)
|
||||
|
||||
# Paths and permissions
|
||||
allowed_paths: List[str] = field(default_factory=list)
|
||||
project_paths: List[str] = field(default_factory=list)
|
||||
project_dir: Optional[str] = None
|
||||
|
||||
# Tool configuration
|
||||
enabled_tools: Dict[str, bool] = field(default_factory=dict)
|
||||
disabled_tools: List[str] = field(default_factory=list)
|
||||
|
||||
# Agent configuration
|
||||
agent: AgentConfig = field(default_factory=AgentConfig)
|
||||
|
||||
# Vector store configuration
|
||||
vector_store: VectorStoreConfig = field(default_factory=VectorStoreConfig)
|
||||
|
||||
# MCP Hub configuration
|
||||
mcp_servers: Dict[str, MCPServerConfig] = field(default_factory=dict)
|
||||
hub_enabled: bool = False
|
||||
trusted_servers: List[str] = field(
|
||||
default_factory=list
|
||||
) # Whitelist of trusted server names
|
||||
|
||||
# Project-specific configurations
|
||||
projects: Dict[str, ProjectConfig] = field(default_factory=dict)
|
||||
current_project: Optional[str] = None
|
||||
|
||||
# Mode configuration
|
||||
active_mode: Optional[str] = None
|
||||
|
||||
def __post_init__(self):
|
||||
"""Initialize default tool states if not specified."""
|
||||
if not self.enabled_tools:
|
||||
self.enabled_tools = {
|
||||
name: config.enabled for name, config in TOOL_REGISTRY.items()
|
||||
}
|
||||
|
||||
# Apply disabled_tools list to enabled_tools dict
|
||||
for tool_name in self.disabled_tools:
|
||||
if tool_name in TOOL_REGISTRY:
|
||||
self.enabled_tools[tool_name] = False
|
||||
|
||||
def is_tool_enabled(self, tool_name: str) -> bool:
|
||||
"""Check if a specific tool is enabled."""
|
||||
# Check disabled_tools list first (highest priority)
|
||||
if tool_name in self.disabled_tools:
|
||||
return False
|
||||
return self.enabled_tools.get(
|
||||
tool_name,
|
||||
TOOL_REGISTRY.get(
|
||||
tool_name, type("obj", (object,), {"enabled": False})
|
||||
).enabled,
|
||||
)
|
||||
|
||||
def enable_tool(self, tool_name: str) -> bool:
|
||||
"""Enable a specific tool."""
|
||||
if tool_name in TOOL_REGISTRY:
|
||||
self.enabled_tools[tool_name] = True
|
||||
# Remove from disabled list if present
|
||||
if tool_name in self.disabled_tools:
|
||||
self.disabled_tools.remove(tool_name)
|
||||
return True
|
||||
return False
|
||||
|
||||
def disable_tool(self, tool_name: str) -> bool:
|
||||
"""Disable a specific tool."""
|
||||
if tool_name in TOOL_REGISTRY:
|
||||
self.enabled_tools[tool_name] = False
|
||||
# Add to disabled list if not present
|
||||
if tool_name not in self.disabled_tools:
|
||||
self.disabled_tools.append(tool_name)
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_enabled_tools(self) -> List[str]:
|
||||
"""Get list of enabled tool names."""
|
||||
return [name for name in TOOL_REGISTRY.keys() if self.is_tool_enabled(name)]
|
||||
|
||||
def get_disabled_tools(self) -> List[str]:
|
||||
"""Get list of disabled tool names."""
|
||||
return [name for name in TOOL_REGISTRY.keys() if not self.is_tool_enabled(name)]
|
||||
|
||||
# MCP Server Management
|
||||
def add_mcp_server(self, server_config: MCPServerConfig) -> bool:
|
||||
"""Add a new MCP server configuration."""
|
||||
if server_config.name in self.mcp_servers:
|
||||
return False # Server already exists
|
||||
self.mcp_servers[server_config.name] = server_config
|
||||
return True
|
||||
|
||||
def remove_mcp_server(self, server_name: str) -> bool:
|
||||
"""Remove an MCP server configuration."""
|
||||
if server_name in self.mcp_servers:
|
||||
del self.mcp_servers[server_name]
|
||||
# Remove from trusted list if present
|
||||
if server_name in self.trusted_servers:
|
||||
self.trusted_servers.remove(server_name)
|
||||
return True
|
||||
return False
|
||||
|
||||
def enable_mcp_server(self, server_name: str) -> bool:
|
||||
"""Enable an MCP server."""
|
||||
if server_name in self.mcp_servers:
|
||||
self.mcp_servers[server_name].enabled = True
|
||||
return True
|
||||
return False
|
||||
|
||||
def disable_mcp_server(self, server_name: str) -> bool:
|
||||
"""Disable an MCP server."""
|
||||
if server_name in self.mcp_servers:
|
||||
self.mcp_servers[server_name].enabled = False
|
||||
return True
|
||||
return False
|
||||
|
||||
def trust_mcp_server(self, server_name: str) -> bool:
|
||||
"""Add server to trusted list."""
|
||||
if server_name in self.mcp_servers:
|
||||
if server_name not in self.trusted_servers:
|
||||
self.trusted_servers.append(server_name)
|
||||
self.mcp_servers[server_name].trusted = True
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_enabled_mcp_servers(self) -> List[MCPServerConfig]:
|
||||
"""Get list of enabled MCP servers."""
|
||||
return [server for server in self.mcp_servers.values() if server.enabled]
|
||||
|
||||
def get_trusted_mcp_servers(self) -> List[MCPServerConfig]:
|
||||
"""Get list of trusted MCP servers."""
|
||||
return [server for server in self.mcp_servers.values() if server.trusted]
|
||||
|
||||
# Project Management
|
||||
def add_project(self, project_config: ProjectConfig) -> bool:
|
||||
"""Add a project configuration."""
|
||||
if project_config.name in self.projects:
|
||||
return False # Project already exists
|
||||
self.projects[project_config.name] = project_config
|
||||
return True
|
||||
|
||||
def remove_project(self, project_name: str) -> bool:
|
||||
"""Remove a project configuration."""
|
||||
if project_name in self.projects:
|
||||
del self.projects[project_name]
|
||||
if self.current_project == project_name:
|
||||
self.current_project = None
|
||||
return True
|
||||
return False
|
||||
|
||||
def set_current_project(self, project_name: str) -> bool:
|
||||
"""Set the current active project."""
|
||||
if project_name in self.projects:
|
||||
self.current_project = project_name
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_current_project(self) -> Optional[ProjectConfig]:
|
||||
"""Get the current project configuration."""
|
||||
if self.current_project:
|
||||
return self.projects.get(self.current_project)
|
||||
return None
|
||||
|
||||
def get_project_tools(self, project_name: Optional[str] = None) -> Dict[str, bool]:
|
||||
"""Get tool configuration for a specific project."""
|
||||
if not project_name:
|
||||
project_name = self.current_project
|
||||
|
||||
if project_name and project_name in self.projects:
|
||||
project = self.projects[project_name]
|
||||
# Start with global tool settings
|
||||
tools = self.enabled_tools.copy()
|
||||
# Apply project-specific settings
|
||||
tools.update(project.enabled_tools)
|
||||
# Apply project-specific disabled tools
|
||||
for tool_name in project.disabled_tools:
|
||||
tools[tool_name] = False
|
||||
return tools
|
||||
|
||||
return self.enabled_tools
|
||||
|
||||
|
||||
def get_config_dir() -> Path:
|
||||
"""Get the configuration directory for Hanzo AI."""
|
||||
if os.name == "nt": # Windows
|
||||
config_dir = Path(os.environ.get("APPDATA", "")) / "hanzo"
|
||||
else: # Unix/macOS
|
||||
config_dir = Path.home() / ".config" / "hanzo"
|
||||
|
||||
config_dir.mkdir(parents=True, exist_ok=True)
|
||||
return config_dir
|
||||
|
||||
|
||||
def get_global_config_path() -> Path:
|
||||
"""Get the path to the global configuration file."""
|
||||
return get_config_dir() / "mcp-settings.json"
|
||||
|
||||
|
||||
def get_project_config_path(project_dir: Optional[str] = None) -> Optional[Path]:
|
||||
"""Get the path to the project-specific configuration file."""
|
||||
if project_dir:
|
||||
project_path = Path(project_dir)
|
||||
config_candidates = [
|
||||
project_path / ".hanzo" / "mcp-settings.json",
|
||||
project_path / "hanzo-mcp.json",
|
||||
project_path / ".hanzo-mcp.json",
|
||||
]
|
||||
for config_path in config_candidates:
|
||||
if config_path.exists():
|
||||
return config_path
|
||||
return None
|
||||
|
||||
|
||||
def ensure_project_hanzo_dir(project_dir: str) -> Path:
|
||||
"""Ensure .hanzo directory exists in project and return its path."""
|
||||
project_path = Path(project_dir)
|
||||
hanzo_dir = project_path / ".hanzo"
|
||||
hanzo_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Create default structure
|
||||
(hanzo_dir / "db").mkdir(exist_ok=True) # Vector database
|
||||
|
||||
# Create default project config if it doesn't exist
|
||||
config_path = hanzo_dir / "mcp-settings.json"
|
||||
if not config_path.exists():
|
||||
default_project_config = {
|
||||
"name": project_path.name,
|
||||
"root_path": str(project_path),
|
||||
"rules": [
|
||||
"Follow project-specific coding standards",
|
||||
"Test all changes before committing",
|
||||
"Update documentation for new features",
|
||||
],
|
||||
"workflows": {
|
||||
"development": {
|
||||
"steps": ["edit", "test", "commit"],
|
||||
"tools": ["read", "write", "edit", "run_command"],
|
||||
},
|
||||
"documentation": {
|
||||
"steps": ["read", "analyze", "write"],
|
||||
"tools": ["read", "write", "vector_search"],
|
||||
},
|
||||
},
|
||||
"tasks": [],
|
||||
"enabled_tools": {},
|
||||
"disabled_tools": [],
|
||||
"mcp_servers": [],
|
||||
}
|
||||
|
||||
with open(config_path, "w") as f:
|
||||
json.dump(default_project_config, f, indent=2)
|
||||
|
||||
return hanzo_dir
|
||||
|
||||
|
||||
def detect_project_from_path(file_path: str) -> Optional[Dict[str, str]]:
|
||||
"""Detect project information from a file path by looking for LLM.md."""
|
||||
path = Path(file_path).resolve()
|
||||
current_path = path.parent if path.is_file() else path
|
||||
|
||||
while current_path != current_path.parent: # Stop at filesystem root
|
||||
llm_md_path = current_path / "LLM.md"
|
||||
if llm_md_path.exists():
|
||||
return {
|
||||
"name": current_path.name,
|
||||
"root_path": str(current_path),
|
||||
"llm_md_path": str(llm_md_path),
|
||||
"hanzo_dir": str(ensure_project_hanzo_dir(str(current_path))),
|
||||
}
|
||||
current_path = current_path.parent
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _load_from_env() -> Dict[str, Any]:
|
||||
"""Load configuration from environment variables."""
|
||||
config = {}
|
||||
|
||||
# Check for agent API keys
|
||||
has_api_keys = False
|
||||
|
||||
# HANZO_API_KEY
|
||||
if hanzo_key := os.environ.get("HANZO_API_KEY"):
|
||||
config.setdefault("agent", {})["hanzo_api_key"] = hanzo_key
|
||||
config["agent"]["enabled"] = True
|
||||
has_api_keys = True
|
||||
|
||||
# Check for other API keys
|
||||
api_key_env_vars = [
|
||||
("OPENAI_API_KEY", "openai"),
|
||||
("ANTHROPIC_API_KEY", "anthropic"),
|
||||
("GOOGLE_API_KEY", "google"),
|
||||
("GROQ_API_KEY", "groq"),
|
||||
("TOGETHER_API_KEY", "together"),
|
||||
("MISTRAL_API_KEY", "mistral"),
|
||||
("PERPLEXITY_API_KEY", "perplexity"),
|
||||
]
|
||||
|
||||
for env_var, _provider in api_key_env_vars:
|
||||
if os.environ.get(env_var):
|
||||
has_api_keys = True
|
||||
break
|
||||
|
||||
# Auto-enable agent and consensus tools if API keys present
|
||||
if has_api_keys:
|
||||
config.setdefault("enabled_tools", {})
|
||||
config["enabled_tools"]["agent"] = True
|
||||
config["enabled_tools"]["consensus"] = True
|
||||
config.setdefault("agent", {})["enabled"] = True
|
||||
|
||||
# Check for MODE/PERSONALITY/HANZO_MODE
|
||||
if (
|
||||
mode := os.environ.get("HANZO_MODE")
|
||||
or os.environ.get("PERSONALITY")
|
||||
or os.environ.get("MODE")
|
||||
):
|
||||
config["active_mode"] = mode
|
||||
|
||||
# Check for other environment overrides
|
||||
if project_dir := os.environ.get("HANZO_PROJECT_DIR"):
|
||||
config["project_dir"] = project_dir
|
||||
|
||||
if log_level := os.environ.get("HANZO_LOG_LEVEL"):
|
||||
config.setdefault("server", {})["log_level"] = log_level
|
||||
|
||||
if allowed_paths := os.environ.get("HANZO_ALLOWED_PATHS"):
|
||||
config["allowed_paths"] = allowed_paths.split(":")
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def load_settings(
|
||||
project_dir: Optional[str] = None, config_overrides: Optional[Dict[str, Any]] = None
|
||||
) -> HanzoMCPSettings:
|
||||
"""Load settings from all sources in priority order.
|
||||
|
||||
Priority (highest to lowest):
|
||||
1. config_overrides (usually from CLI)
|
||||
2. Environment variables
|
||||
3. Project-specific config file
|
||||
4. Global config file
|
||||
5. Defaults
|
||||
"""
|
||||
# Start with defaults
|
||||
settings = HanzoMCPSettings()
|
||||
|
||||
# Load global config
|
||||
global_config_path = get_global_config_path()
|
||||
if global_config_path.exists():
|
||||
try:
|
||||
with open(global_config_path) as f:
|
||||
global_config = json.load(f)
|
||||
settings = _merge_config(settings, global_config)
|
||||
except Exception as e:
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.warning(f"Failed to load global config: {e}")
|
||||
|
||||
# Load project config
|
||||
project_config_path = get_project_config_path(project_dir)
|
||||
if project_config_path:
|
||||
try:
|
||||
with open(project_config_path) as f:
|
||||
project_config = json.load(f)
|
||||
settings = _merge_config(settings, project_config)
|
||||
except Exception as e:
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.warning(f"Failed to load project config: {e}")
|
||||
|
||||
# Apply environment variables
|
||||
env_config = _load_from_env()
|
||||
if env_config:
|
||||
settings = _merge_config(settings, env_config)
|
||||
|
||||
# Apply CLI overrides
|
||||
if config_overrides:
|
||||
settings = _merge_config(settings, config_overrides)
|
||||
|
||||
return settings
|
||||
|
||||
|
||||
def save_settings(settings: HanzoMCPSettings, global_config: bool = True) -> Path:
|
||||
"""Save settings to configuration file.
|
||||
|
||||
Args:
|
||||
settings: Settings to save
|
||||
global_config: If True, save to global config, otherwise save to project config
|
||||
|
||||
Returns:
|
||||
Path where settings were saved
|
||||
"""
|
||||
if global_config:
|
||||
config_path = get_global_config_path()
|
||||
else:
|
||||
# Save to current directory project config
|
||||
config_path = Path.cwd() / ".hanzo-mcp.json"
|
||||
|
||||
config_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(config_path, "w") as f:
|
||||
json.dump(asdict(settings), f, indent=2)
|
||||
|
||||
return config_path
|
||||
|
||||
|
||||
def _merge_config(
|
||||
base_settings: HanzoMCPSettings, config_dict: Dict[str, Any]
|
||||
) -> HanzoMCPSettings:
|
||||
"""Merge configuration dictionary into settings object."""
|
||||
# Convert to dict, merge, then convert back
|
||||
base_dict = asdict(base_settings)
|
||||
|
||||
def deep_merge(base: Dict[str, Any], update: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Deep merge two dictionaries."""
|
||||
for key, value in update.items():
|
||||
if key in base and isinstance(base[key], dict) and isinstance(value, dict):
|
||||
base[key] = deep_merge(base[key], value)
|
||||
else:
|
||||
base[key] = value
|
||||
return base
|
||||
|
||||
merged = deep_merge(base_dict, config_dict)
|
||||
|
||||
# Backwards/forwards compatibility: support a structured "tools" section
|
||||
# where each tool can define { enabled: bool, ...options } and map it to
|
||||
# the existing enabled_tools/disabled_tools layout.
|
||||
tools_cfg = merged.get("tools", {})
|
||||
if isinstance(tools_cfg, dict):
|
||||
enabled_tools = dict(merged.get("enabled_tools", {}))
|
||||
for tool_name, tool_data in tools_cfg.items():
|
||||
if isinstance(tool_data, dict) and "enabled" in tool_data:
|
||||
enabled_tools[tool_name] = bool(tool_data.get("enabled"))
|
||||
merged["enabled_tools"] = enabled_tools
|
||||
|
||||
# Reconstruct the settings object
|
||||
mcp_servers = {}
|
||||
for name, server_data in merged.get("mcp_servers", {}).items():
|
||||
mcp_servers[name] = MCPServerConfig(**server_data)
|
||||
|
||||
projects = {}
|
||||
for name, project_data in merged.get("projects", {}).items():
|
||||
projects[name] = ProjectConfig(**project_data)
|
||||
|
||||
return HanzoMCPSettings(
|
||||
server=ServerConfig(**merged.get("server", {})),
|
||||
allowed_paths=merged.get("allowed_paths", []),
|
||||
project_paths=merged.get("project_paths", []),
|
||||
project_dir=merged.get("project_dir"),
|
||||
enabled_tools=merged.get("enabled_tools", {}),
|
||||
disabled_tools=merged.get("disabled_tools", []),
|
||||
agent=AgentConfig(**merged.get("agent", {})),
|
||||
vector_store=VectorStoreConfig(**merged.get("vector_store", {})),
|
||||
mcp_servers=mcp_servers,
|
||||
hub_enabled=merged.get("hub_enabled", False),
|
||||
trusted_servers=merged.get("trusted_servers", []),
|
||||
projects=projects,
|
||||
current_project=merged.get("current_project"),
|
||||
)
|
||||
@@ -1,209 +0,0 @@
|
||||
"""Tool configuration definitions for Hanzo AI."""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Dict, List
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
class ToolCategory(str, Enum):
|
||||
"""Categories of tools available in Hanzo AI."""
|
||||
|
||||
FILESYSTEM = "filesystem"
|
||||
SHELL = "shell"
|
||||
JUPYTER = "jupyter"
|
||||
TODO = "todo"
|
||||
AGENT = "agent"
|
||||
COMMON = "common"
|
||||
VECTOR = "vector"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolConfig:
|
||||
"""Configuration for an individual tool."""
|
||||
|
||||
name: str
|
||||
category: ToolCategory
|
||||
enabled: bool = True
|
||||
description: str = ""
|
||||
requires_permissions: bool = True
|
||||
cli_flag: str = ""
|
||||
|
||||
def __post_init__(self):
|
||||
"""Generate CLI flag if not provided."""
|
||||
if not self.cli_flag:
|
||||
self.cli_flag = f"--{'enable' if self.enabled else 'disable'}-{self.name.replace('_', '-')}"
|
||||
|
||||
|
||||
# Complete tool registry with all 17 tools
|
||||
TOOL_REGISTRY: Dict[str, ToolConfig] = {
|
||||
# Filesystem Tools (8)
|
||||
"read": ToolConfig(
|
||||
name="read",
|
||||
category=ToolCategory.FILESYSTEM,
|
||||
description="Read file contents with line numbers and truncation support",
|
||||
cli_flag="--disable-read",
|
||||
),
|
||||
"write": ToolConfig(
|
||||
name="write",
|
||||
category=ToolCategory.FILESYSTEM,
|
||||
description="Write content to files, create new files or overwrite existing ones",
|
||||
cli_flag="--disable-write",
|
||||
),
|
||||
"edit": ToolConfig(
|
||||
name="edit",
|
||||
category=ToolCategory.FILESYSTEM,
|
||||
description="Make precise string replacements in files with validation",
|
||||
cli_flag="--disable-edit",
|
||||
),
|
||||
"multi_edit": ToolConfig(
|
||||
name="multi_edit",
|
||||
category=ToolCategory.FILESYSTEM,
|
||||
description="Perform multiple edits to a single file in one operation",
|
||||
cli_flag="--disable-multi-edit",
|
||||
),
|
||||
"directory_tree": ToolConfig(
|
||||
name="directory_tree",
|
||||
category=ToolCategory.FILESYSTEM,
|
||||
description="Display directory structure as a tree",
|
||||
cli_flag="--disable-directory-tree",
|
||||
),
|
||||
"grep": ToolConfig(
|
||||
name="grep",
|
||||
category=ToolCategory.FILESYSTEM,
|
||||
description="Fast content search using ripgrep or fallback Python implementation",
|
||||
cli_flag="--disable-grep",
|
||||
),
|
||||
"ast": ToolConfig(
|
||||
name="ast",
|
||||
category=ToolCategory.FILESYSTEM,
|
||||
description="Search source code with AST context using tree-sitter",
|
||||
cli_flag="--disable-ast",
|
||||
),
|
||||
"content_replace": ToolConfig(
|
||||
name="content_replace",
|
||||
category=ToolCategory.FILESYSTEM,
|
||||
description="Bulk text replacement across multiple files",
|
||||
cli_flag="--disable-content-replace",
|
||||
),
|
||||
# Shell Tools (1)
|
||||
"run_command": ToolConfig(
|
||||
name="run_command",
|
||||
category=ToolCategory.SHELL,
|
||||
description="Execute shell commands with session support",
|
||||
cli_flag="--disable-run-command",
|
||||
),
|
||||
# Jupyter Tools (2)
|
||||
"notebook_read": ToolConfig(
|
||||
name="notebook_read",
|
||||
category=ToolCategory.JUPYTER,
|
||||
description="Read Jupyter notebook files (.ipynb)",
|
||||
cli_flag="--disable-notebook-read",
|
||||
),
|
||||
"notebook_edit": ToolConfig(
|
||||
name="notebook_edit",
|
||||
category=ToolCategory.JUPYTER,
|
||||
description="Edit Jupyter notebook cells (replace, insert, delete)",
|
||||
cli_flag="--disable-notebook-edit",
|
||||
),
|
||||
# Todo Tools (2)
|
||||
"todo_read": ToolConfig(
|
||||
name="todo_read",
|
||||
category=ToolCategory.TODO,
|
||||
description="Read the current todo list for a session",
|
||||
cli_flag="--disable-todo-read",
|
||||
),
|
||||
"todo_write": ToolConfig(
|
||||
name="todo_write",
|
||||
category=ToolCategory.TODO,
|
||||
description="Create and manage structured task lists",
|
||||
cli_flag="--disable-todo-write",
|
||||
),
|
||||
# Agent Tools (3)
|
||||
"dispatch_agent": ToolConfig(
|
||||
name="dispatch_agent",
|
||||
category=ToolCategory.AGENT,
|
||||
enabled=False, # Disabled by default
|
||||
description="Delegate tasks to sub-agents for concurrent/specialized processing",
|
||||
cli_flag="--enable-dispatch-agent",
|
||||
),
|
||||
"swarm": ToolConfig(
|
||||
name="swarm",
|
||||
category=ToolCategory.AGENT,
|
||||
enabled=False, # Disabled by default
|
||||
description="Execute multiple agent tasks in parallel across different files",
|
||||
cli_flag="--enable-swarm",
|
||||
),
|
||||
"hierarchical_swarm": ToolConfig(
|
||||
name="hierarchical_swarm",
|
||||
category=ToolCategory.AGENT,
|
||||
enabled=False, # Disabled by default
|
||||
description="Execute hierarchical agent swarms with Claude Code integration",
|
||||
cli_flag="--enable-hierarchical-swarm",
|
||||
),
|
||||
# Common Tools (3)
|
||||
"think": ToolConfig(
|
||||
name="think",
|
||||
category=ToolCategory.COMMON,
|
||||
description="Provide structured thinking space for complex reasoning",
|
||||
cli_flag="--disable-think",
|
||||
),
|
||||
"batch": ToolConfig(
|
||||
name="batch",
|
||||
category=ToolCategory.COMMON,
|
||||
description="Execute multiple tools in parallel or serial",
|
||||
cli_flag="--disable-batch",
|
||||
),
|
||||
# Vector Tools (2)
|
||||
"vector_index": ToolConfig(
|
||||
name="vector_index",
|
||||
category=ToolCategory.VECTOR,
|
||||
enabled=False, # Disabled by default
|
||||
description="Index documents in local vector database for semantic search",
|
||||
cli_flag="--enable-vector-index",
|
||||
),
|
||||
"vector_search": ToolConfig(
|
||||
name="vector_search",
|
||||
category=ToolCategory.VECTOR,
|
||||
enabled=False, # Disabled by default
|
||||
description="Search documents using semantic similarity in vector database",
|
||||
cli_flag="--enable-vector-search",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def get_tools_by_category(category: ToolCategory) -> List[ToolConfig]:
|
||||
"""Get all tools in a specific category."""
|
||||
return [tool for tool in TOOL_REGISTRY.values() if tool.category == category]
|
||||
|
||||
|
||||
def get_enabled_tools() -> List[ToolConfig]:
|
||||
"""Get all enabled tools."""
|
||||
return [tool for tool in TOOL_REGISTRY.values() if tool.enabled]
|
||||
|
||||
|
||||
def get_disabled_tools() -> List[ToolConfig]:
|
||||
"""Get all disabled tools."""
|
||||
return [tool for tool in TOOL_REGISTRY.values() if not tool.enabled]
|
||||
|
||||
|
||||
def enable_tool(tool_name: str) -> bool:
|
||||
"""Enable a specific tool. Returns True if successful."""
|
||||
if tool_name in TOOL_REGISTRY:
|
||||
TOOL_REGISTRY[tool_name].enabled = True
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def disable_tool(tool_name: str) -> bool:
|
||||
"""Disable a specific tool. Returns True if successful."""
|
||||
if tool_name in TOOL_REGISTRY:
|
||||
TOOL_REGISTRY[tool_name].enabled = False
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_tool_enabled(tool_name: str) -> bool:
|
||||
"""Check if a tool is enabled."""
|
||||
return TOOL_REGISTRY.get(
|
||||
tool_name, ToolConfig("", ToolCategory.COMMON, False)
|
||||
).enabled
|
||||
@@ -1,520 +0,0 @@
|
||||
"""Base Agent - Unified foundation for all AI agent implementations.
|
||||
|
||||
This module provides the single base class for all agent operations,
|
||||
following DRY principles and ensuring consistent behavior across all agents.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import asyncio
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, List, Generic, TypeVar, Optional, Protocol, runtime_checkable
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from dataclasses import field, dataclass
|
||||
|
||||
from .model_registry import ModelConfig, registry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Type variables for generic context
|
||||
TContext = TypeVar("TContext")
|
||||
TResult = TypeVar("TResult")
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class AgentContext(Protocol):
|
||||
"""Protocol for agent execution context."""
|
||||
|
||||
async def log(self, message: str, level: str = "info") -> None:
|
||||
"""Log a message."""
|
||||
pass
|
||||
|
||||
async def progress(self, message: str, percentage: Optional[float] = None) -> None:
|
||||
"""Report progress."""
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentConfig:
|
||||
"""Configuration for agent execution."""
|
||||
|
||||
model: str = "claude-3-5-sonnet-20241022"
|
||||
timeout: int = 300
|
||||
max_retries: int = 3
|
||||
working_dir: Optional[Path] = None
|
||||
environment: Dict[str, str] = field(default_factory=dict)
|
||||
stream_output: bool = False
|
||||
use_worktree: bool = False
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""Resolve model name and validate configuration."""
|
||||
self.model = registry.resolve(self.model)
|
||||
if self.working_dir and not isinstance(self.working_dir, Path):
|
||||
self.working_dir = Path(self.working_dir)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentResult:
|
||||
"""Result from agent execution."""
|
||||
|
||||
success: bool
|
||||
output: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
duration: Optional[float] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def content(self) -> str:
|
||||
"""Get the primary content (output or error)."""
|
||||
return self.output if self.success else (self.error or "Unknown error")
|
||||
|
||||
|
||||
class BaseAgent(ABC, Generic[TContext]):
|
||||
"""Base class for all AI agents.
|
||||
|
||||
This is the single foundation for all agent implementations,
|
||||
ensuring consistent behavior and eliminating code duplication.
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[AgentConfig] = None) -> None:
|
||||
"""Initialize agent with configuration.
|
||||
|
||||
Args:
|
||||
config: Agent configuration
|
||||
"""
|
||||
self.config = config or AgentConfig()
|
||||
self._start_time: Optional[datetime] = None
|
||||
self._end_time: Optional[datetime] = None
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
"""Agent name."""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def description(self) -> str:
|
||||
"""Agent description."""
|
||||
pass
|
||||
|
||||
async def execute(
|
||||
self,
|
||||
prompt: str,
|
||||
context: Optional[TContext] = None,
|
||||
**kwargs: Any,
|
||||
) -> AgentResult:
|
||||
"""Execute agent with prompt.
|
||||
|
||||
Args:
|
||||
prompt: The prompt or task
|
||||
context: Execution context
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
Agent execution result
|
||||
"""
|
||||
self._start_time = datetime.now()
|
||||
|
||||
try:
|
||||
# Setup environment
|
||||
env = self._prepare_environment()
|
||||
|
||||
# Log start
|
||||
if context and isinstance(context, AgentContext):
|
||||
await context.log(f"Starting {self.name} with model {self.config.model}")
|
||||
|
||||
# Execute with retries
|
||||
result = await self._execute_with_retries(prompt, context, env, **kwargs)
|
||||
|
||||
# Calculate duration
|
||||
self._end_time = datetime.now()
|
||||
duration = (self._end_time - self._start_time).total_seconds()
|
||||
|
||||
return AgentResult(
|
||||
success=True,
|
||||
output=result,
|
||||
duration=duration,
|
||||
metadata={"model": self.config.model, "agent": self.name},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self._end_time = datetime.now()
|
||||
duration = (self._end_time - self._start_time).total_seconds() if self._start_time else None
|
||||
|
||||
logger.error(f"Agent {self.name} failed: {e}")
|
||||
|
||||
return AgentResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
duration=duration,
|
||||
metadata={"model": self.config.model, "agent": self.name},
|
||||
)
|
||||
|
||||
def _prepare_environment(self) -> Dict[str, str]:
|
||||
"""Prepare environment variables for execution.
|
||||
|
||||
Returns:
|
||||
Environment variables dictionary
|
||||
"""
|
||||
env = os.environ.copy()
|
||||
|
||||
# Add model-specific API key
|
||||
model_config = registry.get(self.config.model)
|
||||
if model_config and model_config.api_key_env:
|
||||
key_var = model_config.api_key_env
|
||||
if key_var in os.environ:
|
||||
env[key_var] = os.environ[key_var]
|
||||
|
||||
# Add Hanzo unified auth
|
||||
if "HANZO_API_KEY" in os.environ:
|
||||
env["HANZO_API_KEY"] = os.environ["HANZO_API_KEY"]
|
||||
|
||||
# Add custom environment
|
||||
env.update(self.config.environment)
|
||||
|
||||
return env
|
||||
|
||||
async def _execute_with_retries(
|
||||
self,
|
||||
prompt: str,
|
||||
context: Optional[TContext],
|
||||
env: Dict[str, str],
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
"""Execute with retry logic.
|
||||
|
||||
Args:
|
||||
prompt: The prompt
|
||||
context: Execution context
|
||||
env: Environment variables
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
Execution output
|
||||
|
||||
Raises:
|
||||
Exception: If all retries fail
|
||||
"""
|
||||
last_error = None
|
||||
|
||||
for attempt in range(self.config.max_retries):
|
||||
try:
|
||||
# Call the implementation
|
||||
result = await self._execute_impl(prompt, context, env, **kwargs)
|
||||
return result
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
last_error = f"Timeout after {self.config.timeout} seconds"
|
||||
if context and isinstance(context, AgentContext):
|
||||
await context.log(f"Attempt {attempt + 1} timed out", "warning")
|
||||
|
||||
except Exception as e:
|
||||
last_error = str(e)
|
||||
if context and isinstance(context, AgentContext):
|
||||
await context.log(f"Attempt {attempt + 1} failed: {e}", "warning")
|
||||
|
||||
# Don't retry on certain errors
|
||||
if "unauthorized" in str(e).lower() or "forbidden" in str(e).lower():
|
||||
raise
|
||||
|
||||
# Wait before retry (exponential backoff)
|
||||
if attempt < self.config.max_retries - 1:
|
||||
await asyncio.sleep(2 ** attempt)
|
||||
|
||||
raise Exception(f"All {self.config.max_retries} attempts failed. Last error: {last_error}")
|
||||
|
||||
@abstractmethod
|
||||
async def _execute_impl(
|
||||
self,
|
||||
prompt: str,
|
||||
context: Optional[TContext],
|
||||
env: Dict[str, str],
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
"""Implementation-specific execution.
|
||||
|
||||
Args:
|
||||
prompt: The prompt
|
||||
context: Execution context
|
||||
env: Environment variables
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
Execution output
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class CLIAgent(BaseAgent[TContext]):
|
||||
"""Base class for CLI-based agents."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def cli_command(self) -> str:
|
||||
"""CLI command to execute."""
|
||||
pass
|
||||
|
||||
def build_command(self, prompt: str, **kwargs: Any) -> List[str]:
|
||||
"""Build the CLI command.
|
||||
|
||||
Args:
|
||||
prompt: The prompt
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
Command arguments list
|
||||
"""
|
||||
command = [self.cli_command]
|
||||
|
||||
# Add model if specified
|
||||
model_config = registry.get(self.config.model)
|
||||
if model_config:
|
||||
command.extend(["--model", model_config.full_name])
|
||||
|
||||
# Add prompt
|
||||
command.append(prompt)
|
||||
|
||||
return command
|
||||
|
||||
async def _execute_impl(
|
||||
self,
|
||||
prompt: str,
|
||||
context: Optional[TContext],
|
||||
env: Dict[str, str],
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
"""Execute CLI command.
|
||||
|
||||
Args:
|
||||
prompt: The prompt
|
||||
context: Execution context
|
||||
env: Environment variables
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
Command output
|
||||
"""
|
||||
command = self.build_command(prompt, **kwargs)
|
||||
|
||||
# Determine if we need stdin
|
||||
needs_stdin = self.cli_command in ["claude", "cline"]
|
||||
|
||||
# Execute command
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*command,
|
||||
stdin=asyncio.subprocess.PIPE if needs_stdin else None,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
cwd=str(self.config.working_dir) if self.config.working_dir else None,
|
||||
env=env,
|
||||
)
|
||||
|
||||
# Handle timeout
|
||||
try:
|
||||
stdout, stderr = await asyncio.wait_for(
|
||||
process.communicate(prompt.encode() if needs_stdin else None),
|
||||
timeout=self.config.timeout,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
process.kill()
|
||||
raise asyncio.TimeoutError(f"Command timed out after {self.config.timeout} seconds")
|
||||
|
||||
# Check for errors
|
||||
if process.returncode != 0:
|
||||
error_msg = stderr.decode() if stderr else "Command failed"
|
||||
raise Exception(error_msg)
|
||||
|
||||
return stdout.decode()
|
||||
|
||||
|
||||
class APIAgent(BaseAgent[TContext]):
|
||||
"""Base class for API-based agents."""
|
||||
|
||||
async def _execute_impl(
|
||||
self,
|
||||
prompt: str,
|
||||
context: Optional[TContext],
|
||||
env: Dict[str, str],
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
"""Execute via API.
|
||||
|
||||
Args:
|
||||
prompt: The prompt
|
||||
context: Execution context
|
||||
env: Environment variables
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
API response
|
||||
"""
|
||||
# This would be implemented by specific API agents
|
||||
# using the appropriate client library
|
||||
raise NotImplementedError("API agents must implement _execute_impl")
|
||||
|
||||
|
||||
class AgentOrchestrator:
|
||||
"""Orchestrator for managing multiple agents."""
|
||||
|
||||
def __init__(self, default_config: Optional[AgentConfig] = None) -> None:
|
||||
"""Initialize orchestrator.
|
||||
|
||||
Args:
|
||||
default_config: Default configuration for agents
|
||||
"""
|
||||
self.default_config = default_config or AgentConfig()
|
||||
self._agents: Dict[str, BaseAgent] = {}
|
||||
self._semaphore: Optional[asyncio.Semaphore] = None
|
||||
|
||||
def register(self, agent: BaseAgent) -> None:
|
||||
"""Register an agent.
|
||||
|
||||
Args:
|
||||
agent: Agent to register
|
||||
"""
|
||||
self._agents[agent.name] = agent
|
||||
|
||||
def get_agent(self, name: str) -> Optional[BaseAgent]:
|
||||
"""Get agent by name.
|
||||
|
||||
Args:
|
||||
name: Agent name
|
||||
|
||||
Returns:
|
||||
Agent instance or None
|
||||
"""
|
||||
return self._agents.get(name)
|
||||
|
||||
async def execute_single(
|
||||
self,
|
||||
agent_name: str,
|
||||
prompt: str,
|
||||
context: Optional[Any] = None,
|
||||
**kwargs: Any,
|
||||
) -> AgentResult:
|
||||
"""Execute single agent.
|
||||
|
||||
Args:
|
||||
agent_name: Name of agent to use
|
||||
prompt: The prompt
|
||||
context: Execution context
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
Execution result
|
||||
"""
|
||||
agent = self.get_agent(agent_name)
|
||||
if not agent:
|
||||
return AgentResult(
|
||||
success=False,
|
||||
error=f"Agent '{agent_name}' not found",
|
||||
)
|
||||
|
||||
return await agent.execute(prompt, context, **kwargs)
|
||||
|
||||
async def execute_parallel(
|
||||
self,
|
||||
tasks: List[Dict[str, Any]],
|
||||
max_concurrent: int = 5,
|
||||
) -> List[AgentResult]:
|
||||
"""Execute multiple agents in parallel.
|
||||
|
||||
Args:
|
||||
tasks: List of task definitions
|
||||
max_concurrent: Maximum concurrent executions
|
||||
|
||||
Returns:
|
||||
List of results
|
||||
"""
|
||||
self._semaphore = asyncio.Semaphore(max_concurrent)
|
||||
|
||||
async def run_with_semaphore(task: Dict[str, Any]) -> AgentResult:
|
||||
async with self._semaphore:
|
||||
return await self.execute_single(
|
||||
task["agent"],
|
||||
task["prompt"],
|
||||
task.get("context"),
|
||||
**task.get("kwargs", {}),
|
||||
)
|
||||
|
||||
return await asyncio.gather(
|
||||
*[run_with_semaphore(task) for task in tasks],
|
||||
return_exceptions=False,
|
||||
)
|
||||
|
||||
async def execute_consensus(
|
||||
self,
|
||||
prompt: str,
|
||||
agents: List[str],
|
||||
threshold: float = 0.66,
|
||||
) -> Dict[str, Any]:
|
||||
"""Execute consensus operation with multiple agents.
|
||||
|
||||
Args:
|
||||
prompt: The prompt
|
||||
agents: List of agent names
|
||||
threshold: Agreement threshold
|
||||
|
||||
Returns:
|
||||
Consensus results
|
||||
"""
|
||||
# Execute all agents in parallel
|
||||
tasks = [{"agent": agent, "prompt": prompt} for agent in agents]
|
||||
results = await self.execute_parallel(tasks)
|
||||
|
||||
# Analyze consensus
|
||||
successful = [r for r in results if r.success]
|
||||
agreement = len(successful) / len(results) if results else 0
|
||||
|
||||
return {
|
||||
"consensus_reached": agreement >= threshold,
|
||||
"agreement_score": agreement,
|
||||
"individual_results": results,
|
||||
"agents_used": agents,
|
||||
}
|
||||
|
||||
async def execute_chain(
|
||||
self,
|
||||
initial_prompt: str,
|
||||
agents: List[str],
|
||||
) -> List[AgentResult]:
|
||||
"""Execute agents in a chain, passing output forward.
|
||||
|
||||
Args:
|
||||
initial_prompt: Initial prompt
|
||||
agents: List of agent names
|
||||
|
||||
Returns:
|
||||
List of results from each step
|
||||
"""
|
||||
results = []
|
||||
current_prompt = initial_prompt
|
||||
|
||||
for agent_name in agents:
|
||||
result = await self.execute_single(agent_name, current_prompt)
|
||||
results.append(result)
|
||||
|
||||
if result.success and result.output:
|
||||
# Use output as input for next agent
|
||||
current_prompt = f"Review and improve:\n{result.output}"
|
||||
else:
|
||||
# Chain broken
|
||||
break
|
||||
|
||||
return results
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AgentContext",
|
||||
"AgentConfig",
|
||||
"AgentResult",
|
||||
"BaseAgent",
|
||||
"CLIAgent",
|
||||
"APIAgent",
|
||||
"AgentOrchestrator",
|
||||
]
|
||||
@@ -1,436 +0,0 @@
|
||||
"""Unified Model Registry - Single source of truth for all AI model mappings.
|
||||
|
||||
This module provides a centralized registry for AI model configurations,
|
||||
eliminating duplication and ensuring consistency across the codebase.
|
||||
Thread-safe singleton implementation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from enum import Enum
|
||||
from typing import Any, Set, Dict, List, Optional
|
||||
from dataclasses import field, dataclass
|
||||
|
||||
|
||||
class ModelProvider(Enum):
|
||||
"""Enumeration of AI model providers."""
|
||||
|
||||
ANTHROPIC = "anthropic"
|
||||
OPENAI = "openai"
|
||||
GOOGLE = "google"
|
||||
XAI = "xai"
|
||||
OLLAMA = "ollama"
|
||||
DEEPSEEK = "deepseek"
|
||||
MISTRAL = "mistral"
|
||||
META = "meta"
|
||||
HANZO = "hanzo"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelConfig:
|
||||
"""Configuration for a single AI model."""
|
||||
|
||||
full_name: str
|
||||
provider: ModelProvider
|
||||
aliases: Set[str] = field(default_factory=set)
|
||||
default_params: Dict[str, Any] = field(default_factory=dict)
|
||||
supports_vision: bool = False
|
||||
supports_tools: bool = False
|
||||
supports_streaming: bool = True
|
||||
context_window: int = 8192
|
||||
max_output: int = 4096
|
||||
api_key_env: Optional[str] = None
|
||||
cli_command: Optional[str] = None
|
||||
|
||||
|
||||
class ModelRegistry:
|
||||
"""Centralized registry for all AI models.
|
||||
|
||||
Thread-safe singleton implementation ensuring single source of truth
|
||||
for model configurations across the codebase.
|
||||
"""
|
||||
|
||||
_instance: Optional[ModelRegistry] = None
|
||||
_lock = threading.Lock()
|
||||
_models: Dict[str, ModelConfig] = {}
|
||||
_initialized = False
|
||||
|
||||
def __new__(cls) -> ModelRegistry:
|
||||
"""Thread-safe singleton pattern."""
|
||||
if cls._instance is None:
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialize model registry once."""
|
||||
if not self._initialized:
|
||||
with self._lock:
|
||||
if not self._initialized:
|
||||
self._initialize_models()
|
||||
self._initialized = True
|
||||
|
||||
def _initialize_models(self) -> None:
|
||||
"""Initialize all model configurations."""
|
||||
# Claude models
|
||||
self._register(ModelConfig(
|
||||
full_name="claude-3-5-sonnet-20241022",
|
||||
provider=ModelProvider.ANTHROPIC,
|
||||
aliases={"claude", "cc", "claude-code", "sonnet", "sonnet-4.1"},
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
context_window=200000,
|
||||
max_output=8192,
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
cli_command="claude",
|
||||
))
|
||||
|
||||
self._register(ModelConfig(
|
||||
full_name="claude-opus-4-1-20250805",
|
||||
provider=ModelProvider.ANTHROPIC,
|
||||
aliases={"opus", "opus-4.1", "claude-opus"},
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
context_window=200000,
|
||||
max_output=8192,
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
cli_command="claude",
|
||||
))
|
||||
|
||||
self._register(ModelConfig(
|
||||
full_name="claude-3-haiku-20240307",
|
||||
provider=ModelProvider.ANTHROPIC,
|
||||
aliases={"haiku", "claude-haiku"},
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
context_window=200000,
|
||||
max_output=4096,
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
cli_command="claude",
|
||||
))
|
||||
|
||||
# OpenAI models
|
||||
self._register(ModelConfig(
|
||||
full_name="gpt-4-turbo",
|
||||
provider=ModelProvider.OPENAI,
|
||||
aliases={"gpt4", "gpt-4", "codex"},
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
context_window=128000,
|
||||
max_output=4096,
|
||||
api_key_env="OPENAI_API_KEY",
|
||||
cli_command="openai",
|
||||
))
|
||||
|
||||
self._register(ModelConfig(
|
||||
full_name="gpt-5-turbo",
|
||||
provider=ModelProvider.OPENAI,
|
||||
aliases={"gpt5", "gpt-5"},
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
context_window=256000,
|
||||
max_output=16384,
|
||||
api_key_env="OPENAI_API_KEY",
|
||||
cli_command="openai",
|
||||
))
|
||||
|
||||
self._register(ModelConfig(
|
||||
full_name="o1-preview",
|
||||
provider=ModelProvider.OPENAI,
|
||||
aliases={"o1", "openai-o1"},
|
||||
supports_vision=False,
|
||||
supports_tools=False,
|
||||
context_window=128000,
|
||||
max_output=32768,
|
||||
api_key_env="OPENAI_API_KEY",
|
||||
cli_command="openai",
|
||||
))
|
||||
|
||||
# Google models
|
||||
self._register(ModelConfig(
|
||||
full_name="gemini-2.0-flash-exp",
|
||||
provider=ModelProvider.GOOGLE,
|
||||
aliases={"gemini-2", "gemini-2.0", "gemini2"},
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
context_window=1000000,
|
||||
max_output=8192,
|
||||
api_key_env="GEMINI_API_KEY",
|
||||
cli_command="gemini",
|
||||
))
|
||||
|
||||
self._register(ModelConfig(
|
||||
full_name="gemini-exp-1206",
|
||||
provider=ModelProvider.GOOGLE,
|
||||
aliases={"gemini-2.5", "gemini-2.5-pro", "gemini-pro-2.5"},
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
context_window=2000000,
|
||||
max_output=8192,
|
||||
api_key_env="GEMINI_API_KEY",
|
||||
cli_command="gemini",
|
||||
))
|
||||
|
||||
self._register(ModelConfig(
|
||||
full_name="gemini-1.5-pro",
|
||||
provider=ModelProvider.GOOGLE,
|
||||
aliases={"gemini", "gemini-pro", "gemini-1.5"},
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
context_window=2000000,
|
||||
max_output=8192,
|
||||
api_key_env="GEMINI_API_KEY",
|
||||
cli_command="gemini",
|
||||
))
|
||||
|
||||
self._register(ModelConfig(
|
||||
full_name="gemini-1.5-flash",
|
||||
provider=ModelProvider.GOOGLE,
|
||||
aliases={"gemini-flash", "flash"},
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
context_window=1000000,
|
||||
max_output=8192,
|
||||
api_key_env="GEMINI_API_KEY",
|
||||
cli_command="gemini",
|
||||
))
|
||||
|
||||
# xAI models
|
||||
self._register(ModelConfig(
|
||||
full_name="grok-4",
|
||||
provider=ModelProvider.XAI,
|
||||
aliases={"grok", "xai-grok", "grok-2"}, # grok-2 for backward compat
|
||||
supports_vision=True, # Grok-4 supports multimodal
|
||||
supports_tools=True,
|
||||
context_window=128000,
|
||||
max_output=8192,
|
||||
api_key_env="XAI_API_KEY",
|
||||
cli_command="grok",
|
||||
))
|
||||
|
||||
# Ollama models
|
||||
self._register(ModelConfig(
|
||||
full_name="ollama/llama-3.2-3b",
|
||||
provider=ModelProvider.OLLAMA,
|
||||
aliases={"llama", "llama-3.2", "llama3"},
|
||||
supports_vision=False,
|
||||
supports_tools=False,
|
||||
context_window=128000,
|
||||
max_output=4096,
|
||||
api_key_env=None, # Local model
|
||||
cli_command="ollama",
|
||||
))
|
||||
|
||||
self._register(ModelConfig(
|
||||
full_name="ollama/mistral:7b",
|
||||
provider=ModelProvider.MISTRAL,
|
||||
aliases={"mistral", "mistral-7b"},
|
||||
supports_vision=False,
|
||||
supports_tools=False,
|
||||
context_window=32000,
|
||||
max_output=4096,
|
||||
api_key_env=None, # Local model
|
||||
cli_command="ollama",
|
||||
))
|
||||
|
||||
# DeepSeek models
|
||||
self._register(ModelConfig(
|
||||
full_name="deepseek-coder-v2",
|
||||
provider=ModelProvider.DEEPSEEK,
|
||||
aliases={"deepseek", "deepseek-coder"},
|
||||
supports_vision=False,
|
||||
supports_tools=True,
|
||||
context_window=128000,
|
||||
max_output=8192,
|
||||
api_key_env="DEEPSEEK_API_KEY",
|
||||
cli_command="deepseek",
|
||||
))
|
||||
|
||||
def _register(self, config: ModelConfig) -> None:
|
||||
"""Register a model configuration.
|
||||
|
||||
Args:
|
||||
config: Model configuration to register
|
||||
"""
|
||||
# Register by full name
|
||||
self._models[config.full_name] = config
|
||||
|
||||
# Register all aliases
|
||||
for alias in config.aliases:
|
||||
self._models[alias.lower()] = config
|
||||
|
||||
def get(self, model_name: str) -> Optional[ModelConfig]:
|
||||
"""Get model configuration by name or alias.
|
||||
|
||||
Args:
|
||||
model_name: Model name or alias
|
||||
|
||||
Returns:
|
||||
Model configuration or None if not found
|
||||
"""
|
||||
return self._models.get(model_name.lower())
|
||||
|
||||
def resolve(self, model_name: str) -> str:
|
||||
"""Resolve model name or alias to full model name.
|
||||
|
||||
Args:
|
||||
model_name: Model name or alias
|
||||
|
||||
Returns:
|
||||
Full model name, or original if not found
|
||||
"""
|
||||
config = self.get(model_name)
|
||||
return config.full_name if config else model_name
|
||||
|
||||
def get_by_provider(self, provider: ModelProvider) -> List[ModelConfig]:
|
||||
"""Get all unique models for a specific provider.
|
||||
|
||||
Args:
|
||||
provider: Model provider
|
||||
|
||||
Returns:
|
||||
List of unique model configurations
|
||||
"""
|
||||
seen_names = set()
|
||||
results = []
|
||||
for config in self._models.values():
|
||||
if config.provider == provider and config.full_name not in seen_names:
|
||||
seen_names.add(config.full_name)
|
||||
results.append(config)
|
||||
return results
|
||||
|
||||
def get_models_supporting(
|
||||
self,
|
||||
vision: Optional[bool] = None,
|
||||
tools: Optional[bool] = None,
|
||||
streaming: Optional[bool] = None,
|
||||
) -> List[ModelConfig]:
|
||||
"""Get unique models supporting specific features.
|
||||
|
||||
Args:
|
||||
vision: Filter by vision support
|
||||
tools: Filter by tool support
|
||||
streaming: Filter by streaming support
|
||||
|
||||
Returns:
|
||||
List of unique matching model configurations
|
||||
"""
|
||||
seen_names = set()
|
||||
results = []
|
||||
|
||||
for config in self._models.values():
|
||||
if config.full_name in seen_names:
|
||||
continue
|
||||
|
||||
if vision is not None and config.supports_vision != vision:
|
||||
continue
|
||||
if tools is not None and config.supports_tools != tools:
|
||||
continue
|
||||
if streaming is not None and config.supports_streaming != streaming:
|
||||
continue
|
||||
|
||||
seen_names.add(config.full_name)
|
||||
results.append(config)
|
||||
|
||||
return results
|
||||
|
||||
def get_api_key_env(self, model_name: str) -> Optional[str]:
|
||||
"""Get the API key environment variable for a model.
|
||||
|
||||
Args:
|
||||
model_name: Model name or alias
|
||||
|
||||
Returns:
|
||||
Environment variable name or None
|
||||
"""
|
||||
config = self.get(model_name)
|
||||
return config.api_key_env if config else None
|
||||
|
||||
def get_cli_command(self, model_name: str) -> Optional[str]:
|
||||
"""Get the CLI command for a model.
|
||||
|
||||
Args:
|
||||
model_name: Model name or alias
|
||||
|
||||
Returns:
|
||||
CLI command or None
|
||||
"""
|
||||
config = self.get(model_name)
|
||||
return config.cli_command if config else None
|
||||
|
||||
def list_all_models(self) -> List[str]:
|
||||
"""List all unique model full names.
|
||||
|
||||
Returns:
|
||||
Sorted list of unique full model names
|
||||
"""
|
||||
seen_names = set()
|
||||
for config in self._models.values():
|
||||
seen_names.add(config.full_name)
|
||||
return sorted(list(seen_names))
|
||||
|
||||
def list_all_aliases(self) -> Dict[str, str]:
|
||||
"""List all aliases and their full names.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping aliases to full names
|
||||
"""
|
||||
result = {}
|
||||
for key, config in self._models.items():
|
||||
if key != config.full_name:
|
||||
result[key] = config.full_name
|
||||
return result
|
||||
|
||||
|
||||
# Global singleton instance
|
||||
registry = ModelRegistry()
|
||||
|
||||
|
||||
# Convenience functions
|
||||
def resolve_model(model_name: str) -> str:
|
||||
"""Resolve model name or alias to full model name.
|
||||
|
||||
Args:
|
||||
model_name: Model name or alias
|
||||
|
||||
Returns:
|
||||
Full model name
|
||||
"""
|
||||
return registry.resolve(model_name)
|
||||
|
||||
|
||||
def get_model_config(model_name: str) -> Optional[ModelConfig]:
|
||||
"""Get model configuration.
|
||||
|
||||
Args:
|
||||
model_name: Model name or alias
|
||||
|
||||
Returns:
|
||||
Model configuration or None
|
||||
"""
|
||||
return registry.get(model_name)
|
||||
|
||||
|
||||
def get_api_key_env(model_name: str) -> Optional[str]:
|
||||
"""Get API key environment variable for model.
|
||||
|
||||
Args:
|
||||
model_name: Model name or alias
|
||||
|
||||
Returns:
|
||||
Environment variable name or None
|
||||
"""
|
||||
return registry.get_api_key_env(model_name)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ModelProvider",
|
||||
"ModelConfig",
|
||||
"ModelRegistry",
|
||||
"registry",
|
||||
"resolve_model",
|
||||
"get_model_config",
|
||||
"get_api_key_env",
|
||||
]
|
||||
@@ -1,262 +0,0 @@
|
||||
"""Development server with hot reload for Hanzo AI."""
|
||||
|
||||
import time
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Set, Optional
|
||||
from pathlib import Path
|
||||
|
||||
import watchdog.events
|
||||
import watchdog.observers
|
||||
from watchdog.events import FileSystemEventHandler
|
||||
|
||||
from hanzo_mcp.server import HanzoMCPServer
|
||||
|
||||
|
||||
class MCPReloadHandler(FileSystemEventHandler):
|
||||
"""Handler for file system events that triggers MCP server reload."""
|
||||
|
||||
def __init__(self, restart_callback, ignore_patterns: Optional[Set[str]] = None):
|
||||
"""Initialize the reload handler.
|
||||
|
||||
Args:
|
||||
restart_callback: Function to call when files change
|
||||
ignore_patterns: Set of patterns to ignore
|
||||
"""
|
||||
self.restart_callback = restart_callback
|
||||
self.ignore_patterns = ignore_patterns or {
|
||||
"__pycache__",
|
||||
".pyc",
|
||||
".pyo",
|
||||
".git",
|
||||
".pytest_cache",
|
||||
".mypy_cache",
|
||||
".ruff_cache",
|
||||
".coverage",
|
||||
"*.log",
|
||||
".env",
|
||||
".venv",
|
||||
"venv",
|
||||
"node_modules",
|
||||
}
|
||||
self.last_reload = 0
|
||||
self.reload_delay = 0.5 # Debounce delay in seconds
|
||||
|
||||
def should_ignore(self, path: str) -> bool:
|
||||
"""Check if a path should be ignored."""
|
||||
path_obj = Path(path)
|
||||
|
||||
# Check against ignore patterns
|
||||
for pattern in self.ignore_patterns:
|
||||
if pattern in str(path_obj):
|
||||
return True
|
||||
if path_obj.name.endswith(pattern):
|
||||
return True
|
||||
|
||||
# Only watch Python files and config files
|
||||
if path_obj.is_file():
|
||||
allowed_extensions = {".py", ".json", ".yaml", ".yml", ".toml"}
|
||||
if path_obj.suffix not in allowed_extensions:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def on_any_event(self, event):
|
||||
"""Handle any file system event."""
|
||||
if event.is_directory:
|
||||
return
|
||||
|
||||
if self.should_ignore(event.src_path):
|
||||
return
|
||||
|
||||
# Debounce rapid changes
|
||||
current_time = time.time()
|
||||
if current_time - self.last_reload < self.reload_delay:
|
||||
return
|
||||
|
||||
self.last_reload = current_time
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(f"\n🔄 File changed: {event.src_path}")
|
||||
logger.info("🔄 Reloading MCP server...")
|
||||
|
||||
self.restart_callback()
|
||||
|
||||
|
||||
class DevServer:
|
||||
"""Development server with hot reload capability."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str = "hanzo-dev",
|
||||
allowed_paths: Optional[list[str]] = None,
|
||||
project_paths: Optional[list[str]] = None,
|
||||
project_dir: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize the development server.
|
||||
|
||||
Args:
|
||||
name: Server name
|
||||
allowed_paths: Allowed paths for the server
|
||||
project_paths: Project paths
|
||||
project_dir: Project directory
|
||||
**kwargs: Additional arguments for HanzoMCPServer
|
||||
"""
|
||||
self.name = name
|
||||
self.allowed_paths = allowed_paths or []
|
||||
self.project_paths = project_paths
|
||||
self.project_dir = project_dir
|
||||
self.server_kwargs = kwargs
|
||||
self.server_process = None
|
||||
self.observer = None
|
||||
self.running = False
|
||||
|
||||
def create_server(self) -> HanzoMCPServer:
|
||||
"""Create a new MCP server instance."""
|
||||
return HanzoMCPServer(
|
||||
name=self.name,
|
||||
allowed_paths=self.allowed_paths,
|
||||
project_paths=self.project_paths,
|
||||
project_dir=self.project_dir,
|
||||
**self.server_kwargs,
|
||||
)
|
||||
|
||||
def start_file_watcher(self):
|
||||
"""Start watching for file changes."""
|
||||
# Watch the hanzo_mcp package directory
|
||||
package_dir = Path(__file__).parent
|
||||
|
||||
# Create observer and handler
|
||||
self.observer = watchdog.observers.Observer()
|
||||
handler = MCPReloadHandler(self.restart_server)
|
||||
|
||||
# Watch the package directory
|
||||
self.observer.schedule(handler, str(package_dir), recursive=True)
|
||||
|
||||
# Also watch any project directories
|
||||
if self.project_dir:
|
||||
self.observer.schedule(handler, self.project_dir, recursive=True)
|
||||
|
||||
for path in self.allowed_paths:
|
||||
if Path(path).is_dir() and path not in [str(package_dir), self.project_dir]:
|
||||
self.observer.schedule(handler, path, recursive=True)
|
||||
|
||||
self.observer.start()
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(f"👀 Watching for changes in: {package_dir}")
|
||||
if self.project_dir:
|
||||
logger.info(f"👀 Also watching: {self.project_dir}")
|
||||
|
||||
def stop_file_watcher(self):
|
||||
"""Stop the file watcher."""
|
||||
if self.observer and self.observer.is_alive():
|
||||
self.observer.stop()
|
||||
self.observer.join(timeout=2)
|
||||
|
||||
def restart_server(self):
|
||||
"""Restart the MCP server."""
|
||||
# Since MCP servers run in the same process, we need to handle this differently
|
||||
# For now, we'll log a message indicating a restart is needed
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.warning(
|
||||
"\n⚠️ Server restart required. Please restart the MCP client to reload changes."
|
||||
)
|
||||
logger.info(
|
||||
"💡 Tip: In development, consider using the MCP test client for easier reloading."
|
||||
)
|
||||
|
||||
async def run_async(self, transport: str = "stdio"):
|
||||
"""Run the development server asynchronously."""
|
||||
self.running = True
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(f"\n🚀 Starting Hanzo AI in development mode...")
|
||||
|
||||
# Show compute nodes
|
||||
try:
|
||||
from hanzo_mcp.compute_nodes import ComputeNodeDetector
|
||||
|
||||
detector = ComputeNodeDetector()
|
||||
summary = detector.get_node_summary()
|
||||
logger.info(f"🖥️ {summary}")
|
||||
except Exception:
|
||||
# Silently ignore if compute node detection fails
|
||||
pass
|
||||
|
||||
logger.info(f"🔧 Hot reload enabled - watching for file changes")
|
||||
logger.info(f"📁 Project: {self.project_dir or 'current directory'}")
|
||||
logger.info(f"🌐 Transport: {transport}\n")
|
||||
|
||||
# Start file watcher
|
||||
self.start_file_watcher()
|
||||
|
||||
try:
|
||||
# Create and run server
|
||||
server = self.create_server()
|
||||
|
||||
# Run the server (this will block)
|
||||
server.run(transport=transport)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
logger.info("\n\n🛑 Shutting down development server...")
|
||||
finally:
|
||||
self.running = False
|
||||
self.stop_file_watcher()
|
||||
logger.info("👋 Development server stopped")
|
||||
|
||||
def run(self, transport: str = "stdio"):
|
||||
"""Run the development server."""
|
||||
try:
|
||||
# Run the async version
|
||||
asyncio.run(self.run_async(transport))
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
|
||||
|
||||
def run_dev_server():
|
||||
"""Entry point for development server."""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Run Hanzo AI in development mode with hot reload"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--name", type=str, default="hanzo-dev", help="Name of the MCP server"
|
||||
)
|
||||
parser.add_argument("--project-dir", type=str, help="Project directory to serve")
|
||||
parser.add_argument(
|
||||
"--allowed-path",
|
||||
type=str,
|
||||
action="append",
|
||||
dest="allowed_paths",
|
||||
help="Additional allowed paths (can be specified multiple times)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--transport",
|
||||
type=str,
|
||||
default="stdio",
|
||||
choices=["stdio", "sse"],
|
||||
help="Transport type (default: stdio)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--host", type=str, default="127.0.0.1", help="Host for SSE transport"
|
||||
)
|
||||
parser.add_argument("--port", type=int, default=3000, help="Port for SSE transport")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Create and run dev server
|
||||
dev_server = DevServer(
|
||||
name=args.name,
|
||||
allowed_paths=args.allowed_paths,
|
||||
project_dir=args.project_dir,
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
)
|
||||
|
||||
dev_server.run(transport=args.transport)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_dev_server()
|
||||
@@ -1,242 +0,0 @@
|
||||
import os
|
||||
|
||||
from mcp.server import FastMCP
|
||||
|
||||
from hanzo_mcp.prompts.utils import (
|
||||
get_os_info,
|
||||
get_git_info,
|
||||
get_directory_structure,
|
||||
)
|
||||
from hanzo_mcp.prompts.tool_explorer import (
|
||||
AGENT_TOOLS_HELP,
|
||||
SHELL_TOOLS_HELP,
|
||||
BATCH_TOOL_EXAMPLES,
|
||||
TOOL_EXPLORER_PROMPT,
|
||||
FILESYSTEM_TOOLS_HELP,
|
||||
)
|
||||
from hanzo_mcp.prompts.create_release import CREATE_RELEASE_PROMPT
|
||||
from hanzo_mcp.prompts.project_system import PROJECT_SYSTEM_PROMPT
|
||||
from hanzo_mcp.prompts.enhanced_prompts import (
|
||||
QUICK_START_PROMPT,
|
||||
LSP_TOOLS_HELP_PROMPT,
|
||||
PAGINATION_GUIDE_PROMPT,
|
||||
PERFORMANCE_TIPS_PROMPT,
|
||||
MEMORY_VECTOR_HELP_PROMPT,
|
||||
CONFIGURATION_GUIDE_PROMPT,
|
||||
DATABASE_TOOLS_HELP_PROMPT,
|
||||
NETWORK_AGENT_GUIDE_PROMPT,
|
||||
TROUBLESHOOTING_GUIDE_PROMPT,
|
||||
SECURITY_BEST_PRACTICES_PROMPT,
|
||||
)
|
||||
from hanzo_mcp.prompts.compact_conversation import COMPACT_CONVERSATION_PROMPT
|
||||
from hanzo_mcp.prompts.project_todo_reminder import (
|
||||
PROJECT_TODO_EMPTY_REMINDER,
|
||||
get_project_todo_reminder,
|
||||
)
|
||||
|
||||
CONTINUE_FROM_LAST_SESSION_PROMPT = """<system-reminder>
|
||||
This is a reminder that your todo list is currently empty. DO NOT mention this to the user explicitly because they are already aware. If you are working on tasks that would benefit from a todo list please use the TodoWrite tool to create one. If not, please feel free to ignore. Again do not mention this message to the user.
|
||||
</system-reminder>
|
||||
"""
|
||||
|
||||
|
||||
def create_project_system_prompt(project_path: str):
|
||||
"""Factory function to create a project system prompt function."""
|
||||
|
||||
def project_system_prompt() -> str:
|
||||
"""
|
||||
Summarize the conversation so far for a specific project.
|
||||
"""
|
||||
working_directory = project_path
|
||||
is_git_repo = os.path.isdir(os.path.join(working_directory, ".git"))
|
||||
platform, _, os_version = get_os_info()
|
||||
|
||||
# Get directory structure
|
||||
directory_structure = get_directory_structure(
|
||||
working_directory, max_depth=3, include_filtered=False
|
||||
)
|
||||
|
||||
# Get git information
|
||||
git_info = get_git_info(working_directory)
|
||||
current_branch = git_info.get("current_branch", "")
|
||||
main_branch = git_info.get("main_branch", "")
|
||||
git_status = git_info.get("git_status", "")
|
||||
recent_commits = git_info.get("recent_commits", "")
|
||||
|
||||
return PROJECT_SYSTEM_PROMPT.format(
|
||||
working_directory=working_directory,
|
||||
is_git_repo=is_git_repo,
|
||||
platform=platform,
|
||||
os_version=os_version,
|
||||
directory_structure=directory_structure,
|
||||
current_branch=current_branch,
|
||||
main_branch=main_branch,
|
||||
git_status=git_status,
|
||||
recent_commits=recent_commits,
|
||||
)
|
||||
|
||||
return project_system_prompt
|
||||
|
||||
|
||||
def register_all_prompts(
|
||||
mcp_server: FastMCP, projects: list[str] | None = None
|
||||
) -> None:
|
||||
@mcp_server.prompt(name="Compact current conversation")
|
||||
def compact() -> str:
|
||||
"""
|
||||
Summarize the conversation so far.
|
||||
"""
|
||||
return COMPACT_CONVERSATION_PROMPT
|
||||
|
||||
@mcp_server.prompt(name="Create a new release")
|
||||
def create_release() -> str:
|
||||
"""
|
||||
Create a new release for my project.
|
||||
"""
|
||||
return CREATE_RELEASE_PROMPT
|
||||
|
||||
@mcp_server.prompt(name="Continue todo by session id")
|
||||
def continue_todo_by_session_id(session_id: str) -> str:
|
||||
"""
|
||||
Continue from the last todo list for the current session.
|
||||
"""
|
||||
return get_project_todo_reminder(session_id)
|
||||
|
||||
@mcp_server.prompt(name="Continue latest todo")
|
||||
def continue_latest_todo() -> str:
|
||||
"""
|
||||
Continue from the last todo list for the current session.
|
||||
"""
|
||||
return get_project_todo_reminder()
|
||||
|
||||
@mcp_server.prompt(name="System prompt")
|
||||
def manual_project_system_prompt(project_path: str) -> str:
|
||||
"""
|
||||
Detailed system prompt include env,git etc information about the specified project.
|
||||
"""
|
||||
return create_project_system_prompt(project_path)()
|
||||
|
||||
@mcp_server.prompt(name="Explore all tools")
|
||||
def explore_tools() -> str:
|
||||
"""
|
||||
Comprehensive guide to all available Hanzo MCP tools and how to use them.
|
||||
"""
|
||||
return TOOL_EXPLORER_PROMPT
|
||||
|
||||
@mcp_server.prompt(name="Filesystem tools help")
|
||||
def filesystem_help() -> str:
|
||||
"""
|
||||
Detailed guide for filesystem tools (read, write, edit, search, etc).
|
||||
"""
|
||||
return FILESYSTEM_TOOLS_HELP
|
||||
|
||||
@mcp_server.prompt(name="Agent tools help")
|
||||
def agent_help() -> str:
|
||||
"""
|
||||
Guide for using agent tools to delegate complex tasks.
|
||||
"""
|
||||
return AGENT_TOOLS_HELP
|
||||
|
||||
@mcp_server.prompt(name="Shell tools help")
|
||||
def shell_help() -> str:
|
||||
"""
|
||||
Guide for shell and command execution tools.
|
||||
"""
|
||||
return SHELL_TOOLS_HELP
|
||||
|
||||
@mcp_server.prompt(name="Batch tool examples")
|
||||
def batch_examples() -> str:
|
||||
"""
|
||||
Advanced examples of using the batch tool for parallel operations.
|
||||
"""
|
||||
return BATCH_TOOL_EXAMPLES
|
||||
|
||||
# Enhanced prompts for better discoverability
|
||||
@mcp_server.prompt(name="Quick start guide")
|
||||
def quick_start() -> str:
|
||||
"""
|
||||
Common workflows and recipes for getting started quickly.
|
||||
"""
|
||||
return QUICK_START_PROMPT
|
||||
|
||||
@mcp_server.prompt(name="Pagination guide")
|
||||
def pagination_guide() -> str:
|
||||
"""
|
||||
How to use pagination for large result sets.
|
||||
"""
|
||||
return PAGINATION_GUIDE_PROMPT
|
||||
|
||||
@mcp_server.prompt(name="Memory and vector tools help")
|
||||
def memory_vector_help() -> str:
|
||||
"""
|
||||
Guide for semantic search and memory tools.
|
||||
"""
|
||||
return MEMORY_VECTOR_HELP_PROMPT
|
||||
|
||||
@mcp_server.prompt(name="Database tools help")
|
||||
def database_help() -> str:
|
||||
"""
|
||||
SQL and graph database operations guide.
|
||||
"""
|
||||
return DATABASE_TOOLS_HELP_PROMPT
|
||||
|
||||
@mcp_server.prompt(name="LSP tools help")
|
||||
def lsp_help() -> str:
|
||||
"""
|
||||
Language Server Protocol features and code intelligence.
|
||||
"""
|
||||
return LSP_TOOLS_HELP_PROMPT
|
||||
|
||||
@mcp_server.prompt(name="Configuration guide")
|
||||
def config_guide() -> str:
|
||||
"""
|
||||
How to configure tools, presets, and settings.
|
||||
"""
|
||||
return CONFIGURATION_GUIDE_PROMPT
|
||||
|
||||
@mcp_server.prompt(name="Network agent guide")
|
||||
def network_guide() -> str:
|
||||
"""
|
||||
Distributed AI orchestration with network/swarm tools.
|
||||
"""
|
||||
return NETWORK_AGENT_GUIDE_PROMPT
|
||||
|
||||
@mcp_server.prompt(name="Performance tips")
|
||||
def performance_tips() -> str:
|
||||
"""
|
||||
Optimization strategies for better performance.
|
||||
"""
|
||||
return PERFORMANCE_TIPS_PROMPT
|
||||
|
||||
@mcp_server.prompt(name="Security best practices")
|
||||
def security_practices() -> str:
|
||||
"""
|
||||
Safe usage patterns and security guidelines.
|
||||
"""
|
||||
return SECURITY_BEST_PRACTICES_PROMPT
|
||||
|
||||
@mcp_server.prompt(name="Troubleshooting guide")
|
||||
def troubleshooting() -> str:
|
||||
"""
|
||||
Common issues and their solutions.
|
||||
"""
|
||||
return TROUBLESHOOTING_GUIDE_PROMPT
|
||||
|
||||
if projects is None:
|
||||
return
|
||||
|
||||
for project in projects:
|
||||
# Register the prompt with the factory function
|
||||
mcp_server.prompt(
|
||||
name=f"System prompt for {os.path.basename(project)}",
|
||||
description=f"Detailed system prompt include env,git etc information about {project}",
|
||||
)(create_project_system_prompt(project))
|
||||
|
||||
return
|
||||
|
||||
|
||||
__all__ = [
|
||||
"register_all_prompts",
|
||||
"get_project_todo_reminder",
|
||||
"PROJECT_TODO_EMPTY_REMINDER",
|
||||
]
|
||||
@@ -1,77 +0,0 @@
|
||||
COMPACT_CONVERSATION_PROMPT = """Your task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions.
|
||||
This summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing development work without losing context.
|
||||
|
||||
Before providing your final summary, wrap your analysis in <analysis> tags to organize your thoughts and ensure you've covered all necessary points. In your analysis process:
|
||||
|
||||
1. Chronologically analyze each message and section of the conversation. For each section thoroughly identify:
|
||||
- The user's explicit requests and intents
|
||||
- Your approach to addressing the user's requests
|
||||
- Key decisions, technical concepts and code patterns
|
||||
- Specific details like file names, full code snippets, function signatures, file edits, etc
|
||||
2. Double-check for technical accuracy and completeness, addressing each required element thoroughly.
|
||||
|
||||
Your summary should include the following sections:
|
||||
|
||||
1. Primary Request and Intent: Capture all of the user's explicit requests and intents in detail
|
||||
2. Key Technical Concepts: List all important technical concepts, technologies, and frameworks discussed.
|
||||
3. Files and Code Sections: Enumerate specific files and code sections examined, modified, or created. Pay special attention to the most recent messages and include full code snippets where applicable and include a summary of why this file read or edit is important.
|
||||
4. Problem Solving: Document problems solved and any ongoing troubleshooting efforts.
|
||||
5. Pending Tasks: Outline any pending tasks that you have explicitly been asked to work on.
|
||||
6. Current Work: Describe in detail precisely what was being worked on immediately before this summary request, paying special attention to the most recent messages from both user and assistant. Include file names and code snippets where applicable.
|
||||
7. Optional Next Step: List the next step that you will take that is related to the most recent work you were doing. IMPORTANT: ensure that this step is DIRECTLY in line with the user's explicit requests, and the task you were working on immediately before this summary request. If your last task was concluded, then only list next steps if they are explicitly in line with the users request. Do not start on tangential requests without confirming with the user first.
|
||||
If there is a next step, include direct quotes from the most recent conversation showing exactly what task you were working on and where you left off. This should be verbatim to ensure there's no drift in task interpretation.
|
||||
|
||||
Here's an example of how your output should be structured:
|
||||
|
||||
<example>
|
||||
<analysis>
|
||||
[Your thought process, ensuring all points are covered thoroughly and accurately]
|
||||
</analysis>
|
||||
|
||||
<summary>
|
||||
1. Primary Request and Intent:
|
||||
[Detailed description]
|
||||
|
||||
2. Key Technical Concepts:
|
||||
- [Concept 1]
|
||||
- [Concept 2]
|
||||
- [...]
|
||||
|
||||
3. Files and Code Sections:
|
||||
- [File Name 1]
|
||||
- [Summary of why this file is important]
|
||||
- [Summary of the changes made to this file, if any]
|
||||
- [Important Code Snippet]
|
||||
- [File Name 2]
|
||||
- [Important Code Snippet]
|
||||
- [...]
|
||||
|
||||
4. Problem Solving:
|
||||
[Description of solved problems and ongoing troubleshooting]
|
||||
|
||||
5. Pending Tasks:
|
||||
- [Task 1]
|
||||
- [Task 2]
|
||||
- [...]
|
||||
|
||||
6. Current Work:
|
||||
[Precise description of current work]
|
||||
|
||||
7. Optional Next Step:
|
||||
[Optional Next step to take]
|
||||
|
||||
</summary>
|
||||
</example>
|
||||
|
||||
Please provide your summary based on the conversation so far, following this structure and ensuring precision and thoroughness in your response.
|
||||
|
||||
There may be additional summarization instructions provided in the included context. If so, remember to follow these instructions when creating the above summary. Examples of instructions include:
|
||||
<example>
|
||||
## Compact Instructions
|
||||
When summarizing the conversation focus on typescript code changes and also remember the mistakes you made and how you fixed them.
|
||||
</example>
|
||||
|
||||
<example>
|
||||
# Summary instructions
|
||||
When you are using compact - please focus on test output and code changes. Include file reads verbatim.
|
||||
</example>"""
|
||||
@@ -1,38 +0,0 @@
|
||||
CREATE_RELEASE_PROMPT = """Help me create a new release for my project. Follow these steps:
|
||||
|
||||
## Initial Analysis
|
||||
1. Examine the project version files (typically `__init__.py`, `package.json`, `pyproject.toml`, etc.)
|
||||
2. Review the current CHANGELOG.md format and previous releases
|
||||
3. Check the release workflow configuration (GitHub Actions, CI/CD pipelines)
|
||||
4. Review commits since the last release tag:
|
||||
```bash
|
||||
git log <last-tag>..HEAD --pretty=format:"%h %s%n%b" --name-status
|
||||
```
|
||||
|
||||
## Version Update
|
||||
1. Identify all files containing version numbers
|
||||
2. Update version numbers consistently across all files
|
||||
3. Follow semantic versioning guidelines (MAJOR.MINOR.PATCH)
|
||||
|
||||
## Changelog Creation
|
||||
1. Add a new section at the top of CHANGELOG.md with the new version and today's date
|
||||
2. Group changes by type: Added, Changed, Fixed, Removed, etc.
|
||||
3. Include commit hashes in parentheses for reference
|
||||
4. Write clear, detailed descriptions for each change
|
||||
5. Follow established project conventions for changelog format
|
||||
|
||||
## Release Commit and Tag
|
||||
1. Commit the version and changelog updates:
|
||||
```bash
|
||||
git add <changed-files>
|
||||
git commit -m "chore: bump version to X.Y.Z"
|
||||
```
|
||||
2. Create an annotated tag:
|
||||
```bash
|
||||
git tag -a "vX.Y.Z" -m "Release vX.Y.Z"
|
||||
```
|
||||
3. Push the changes and tag:
|
||||
```bash
|
||||
git push origin main
|
||||
git push origin vX.Y.Z
|
||||
```"""
|
||||
@@ -1,809 +0,0 @@
|
||||
"""Enhanced prompts for better discoverability and usability."""
|
||||
|
||||
QUICK_START_PROMPT = """# Hanzo MCP Quick Start Guide
|
||||
|
||||
## Common Workflows
|
||||
|
||||
### 1. Explore a New Codebase
|
||||
```python
|
||||
# Get project overview
|
||||
directory_tree(path=".", depth=3)
|
||||
|
||||
# Find main entry points
|
||||
find(pattern="main.*|index.*", path=".")
|
||||
|
||||
# Search for key patterns
|
||||
search(pattern="TODO|FIXME|BUG", path=".")
|
||||
```
|
||||
|
||||
### 2. Multi-Agent Code Analysis
|
||||
```python
|
||||
# Run parallel analysis with different agents
|
||||
batch(
|
||||
description="Comprehensive code analysis",
|
||||
invocations=[
|
||||
{"tool_name": "agent", "input": {"prompt": "Analyze security vulnerabilities"}},
|
||||
{"tool_name": "agent", "input": {"prompt": "Find performance bottlenecks"}},
|
||||
{"tool_name": "agent", "input": {"prompt": "Review code quality and suggest improvements"}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### 3. Refactor Code Across Files
|
||||
```python
|
||||
# Find all occurrences
|
||||
search(pattern="oldFunction", path="./src")
|
||||
|
||||
# Review with critic
|
||||
critic(analysis="Review the usage of oldFunction and suggest refactoring approach")
|
||||
|
||||
# Make changes
|
||||
content_replace(
|
||||
pattern="oldFunction",
|
||||
replacement="newFunction",
|
||||
path="./src"
|
||||
)
|
||||
|
||||
# Run tests
|
||||
bash("npm test")
|
||||
```
|
||||
|
||||
### 4. Track Complex Tasks
|
||||
```python
|
||||
# Create todo list
|
||||
todo_write(todos=[
|
||||
{"content": "Analyze current implementation", "status": "pending"},
|
||||
{"content": "Design new architecture", "status": "pending"},
|
||||
{"content": "Implement changes", "status": "pending"},
|
||||
{"content": "Write tests", "status": "pending"},
|
||||
{"content": "Update documentation", "status": "pending"}
|
||||
])
|
||||
|
||||
# Update as you work
|
||||
todo(action="update", id="task-1", status="completed")
|
||||
```
|
||||
|
||||
## Tips
|
||||
- Use `batch()` for parallel operations
|
||||
- Use `think()` for complex reasoning
|
||||
- Use `critic()` for code review
|
||||
- Use pagination for large result sets"""
|
||||
|
||||
PAGINATION_GUIDE_PROMPT = """# Pagination Guide
|
||||
|
||||
## How Pagination Works
|
||||
|
||||
Most tools support pagination to handle large result sets efficiently.
|
||||
|
||||
### Basic Pagination Parameters
|
||||
- `page`: Page number (starts at 1)
|
||||
- `page_size`: Results per page (default: 50-100)
|
||||
- `max_results`: Total maximum results
|
||||
|
||||
### Examples
|
||||
|
||||
#### Find Tool with Pagination
|
||||
```python
|
||||
# Get first page of Python files
|
||||
find(pattern="*.py", path="/project", page=1, page_size=10)
|
||||
|
||||
# Response includes:
|
||||
{
|
||||
"results": [...],
|
||||
"pagination": {
|
||||
"page": 1,
|
||||
"page_size": 10,
|
||||
"total_results": 150,
|
||||
"total_pages": 15,
|
||||
"has_next": true,
|
||||
"has_prev": false
|
||||
}
|
||||
}
|
||||
|
||||
# Get next page
|
||||
find(pattern="*.py", path="/project", page=2, page_size=10)
|
||||
```
|
||||
|
||||
#### Search with Limits
|
||||
```python
|
||||
# Limit total results
|
||||
search(pattern="TODO", path="/project", max_results=20)
|
||||
```
|
||||
|
||||
### Batch Processing Pages
|
||||
```python
|
||||
# Process multiple pages in parallel
|
||||
batch(
|
||||
description="Get all Python files",
|
||||
invocations=[
|
||||
{"tool_name": "find", "input": {"pattern": "*.py", "page": 1, "page_size": 50}},
|
||||
{"tool_name": "find", "input": {"pattern": "*.py", "page": 2, "page_size": 50}},
|
||||
{"tool_name": "find", "input": {"pattern": "*.py", "page": 3, "page_size": 50}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
## Best Practices
|
||||
1. Start with smaller page sizes for testing
|
||||
2. Use `max_results` to limit total processing
|
||||
3. Check `has_next` before requesting next page
|
||||
4. Use batch for parallel page fetching"""
|
||||
|
||||
MEMORY_VECTOR_HELP_PROMPT = """# Memory & Vector Tools Guide
|
||||
|
||||
## Semantic Search and Memory
|
||||
|
||||
### Vector Search
|
||||
Search using semantic similarity rather than exact text matching.
|
||||
|
||||
```python
|
||||
# Enable vector search in unified search
|
||||
search(
|
||||
pattern="how to handle user authentication",
|
||||
path="/project",
|
||||
enable_vector=true
|
||||
)
|
||||
|
||||
# Index project for vector search
|
||||
vector_index(path="/project", recursive=true)
|
||||
|
||||
# Pure vector search
|
||||
vector_search(
|
||||
query="error handling patterns",
|
||||
path="/project",
|
||||
top_k=10
|
||||
)
|
||||
```
|
||||
|
||||
### Memory Tools
|
||||
Store and retrieve context across sessions.
|
||||
|
||||
```python
|
||||
# Store knowledge
|
||||
memory_add(
|
||||
key="project_architecture",
|
||||
content="The project uses a microservices architecture with..."
|
||||
)
|
||||
|
||||
# Retrieve knowledge
|
||||
memory_get(key="project_architecture")
|
||||
|
||||
# Search memory
|
||||
memory_search(query="architecture decisions")
|
||||
```
|
||||
|
||||
## Use Cases
|
||||
1. **Semantic Code Search**: Find conceptually similar code
|
||||
2. **Knowledge Base**: Store project understanding
|
||||
3. **Context Preservation**: Maintain context across sessions
|
||||
4. **Pattern Discovery**: Find similar implementations
|
||||
|
||||
## Tips
|
||||
- Index large projects once, search many times
|
||||
- Combine vector search with traditional search
|
||||
- Use memory for important discoveries
|
||||
- Vector search works best with natural language queries"""
|
||||
|
||||
DATABASE_TOOLS_HELP_PROMPT = """# Database Tools Guide
|
||||
|
||||
## SQL Database Operations
|
||||
|
||||
### Query Execution
|
||||
```python
|
||||
# Execute SQL query
|
||||
sql_query(
|
||||
query="SELECT * FROM users WHERE created_at > '2024-01-01'",
|
||||
database="myapp.db"
|
||||
)
|
||||
|
||||
# Search across tables
|
||||
sql_search(
|
||||
pattern="[email protected]",
|
||||
database="myapp.db"
|
||||
)
|
||||
|
||||
# Get database statistics
|
||||
sql_stats(database="myapp.db")
|
||||
```
|
||||
|
||||
## Graph Database Operations
|
||||
|
||||
### Managing Graph Data
|
||||
```python
|
||||
# Add nodes and edges
|
||||
graph_add(
|
||||
node_type="User",
|
||||
node_id="user_123",
|
||||
properties={"name": "John", "email": "[email protected]"}
|
||||
)
|
||||
|
||||
graph_add(
|
||||
edge_type="FOLLOWS",
|
||||
from_node="user_123",
|
||||
to_node="user_456"
|
||||
)
|
||||
|
||||
# Query graph
|
||||
graph_query(
|
||||
query="MATCH (u:User)-[:FOLLOWS]->(f:User) RETURN u, f",
|
||||
database="social.graph"
|
||||
)
|
||||
|
||||
# Search graph
|
||||
graph_search(
|
||||
pattern="John",
|
||||
node_type="User"
|
||||
)
|
||||
|
||||
# Graph statistics
|
||||
graph_stats(database="social.graph")
|
||||
```
|
||||
|
||||
## Best Practices
|
||||
1. Always sanitize inputs to prevent injection
|
||||
2. Use transactions for multiple operations
|
||||
3. Index frequently queried fields
|
||||
4. Monitor query performance with stats tools"""
|
||||
|
||||
LSP_TOOLS_HELP_PROMPT = """# Language Server Protocol (LSP) Tools Guide
|
||||
|
||||
## Code Intelligence Features
|
||||
|
||||
### Basic Operations
|
||||
```python
|
||||
# Go to definition
|
||||
lsp(
|
||||
action="definition",
|
||||
file="/src/main.py",
|
||||
line=42,
|
||||
character=15
|
||||
)
|
||||
|
||||
# Find all references
|
||||
lsp(
|
||||
action="references",
|
||||
file="/src/main.py",
|
||||
line=42,
|
||||
character=15
|
||||
)
|
||||
|
||||
# Get hover information
|
||||
lsp(
|
||||
action="hover",
|
||||
file="/src/main.py",
|
||||
line=42,
|
||||
character=15
|
||||
)
|
||||
```
|
||||
|
||||
### Refactoring
|
||||
```python
|
||||
# Rename symbol across codebase
|
||||
lsp(
|
||||
action="rename",
|
||||
file="/src/main.py",
|
||||
line=42,
|
||||
character=15,
|
||||
new_name="newFunctionName"
|
||||
)
|
||||
```
|
||||
|
||||
### Diagnostics
|
||||
```python
|
||||
# Get errors and warnings
|
||||
lsp(
|
||||
action="diagnostics",
|
||||
file="/src/main.py"
|
||||
)
|
||||
```
|
||||
|
||||
### Code Completion
|
||||
```python
|
||||
# Get completions at position
|
||||
lsp(
|
||||
action="completion",
|
||||
file="/src/main.py",
|
||||
line=42,
|
||||
character=15
|
||||
)
|
||||
```
|
||||
|
||||
## Supported Languages
|
||||
- Python (pylsp)
|
||||
- TypeScript/JavaScript (typescript-language-server)
|
||||
- Go (gopls)
|
||||
- Rust (rust-analyzer)
|
||||
- Java (jdtls)
|
||||
- C/C++ (clangd)
|
||||
|
||||
## Tips
|
||||
- LSP servers are installed automatically
|
||||
- Use for accurate refactoring
|
||||
- Combine with search tools for comprehensive analysis"""
|
||||
|
||||
CONFIGURATION_GUIDE_PROMPT = """# Configuration Guide
|
||||
|
||||
## Tool Configuration
|
||||
|
||||
### View Current Configuration
|
||||
```python
|
||||
# List all enabled tools
|
||||
tool_list()
|
||||
|
||||
# Check specific tool status
|
||||
stats()
|
||||
```
|
||||
|
||||
### Enable/Disable Tools
|
||||
```python
|
||||
# Enable a tool
|
||||
tool_enable(name="lsp")
|
||||
|
||||
# Disable a tool
|
||||
tool_disable(name="sql_query")
|
||||
```
|
||||
|
||||
## Configuration Presets
|
||||
|
||||
### Available Presets
|
||||
|
||||
1. **minimal** - Essential tools only
|
||||
- Basic file operations
|
||||
- Simple commands
|
||||
|
||||
2. **standard** - Common development tools
|
||||
- File operations
|
||||
- Search tools
|
||||
- Process management
|
||||
|
||||
3. **development** - Full development suite
|
||||
- All standard tools
|
||||
- Agent tools
|
||||
- Package managers
|
||||
- LSP support
|
||||
|
||||
4. **full** - Everything enabled
|
||||
- All tools available
|
||||
- Maximum capabilities
|
||||
|
||||
5. **ai_research** - AI/ML focused
|
||||
- Agent orchestration
|
||||
- Vector search
|
||||
- Memory tools
|
||||
|
||||
### Switching Presets
|
||||
Configure in your launch command:
|
||||
```bash
|
||||
# Use development preset
|
||||
hanzo-mcp --preset development
|
||||
|
||||
# Or set in config file
|
||||
~/.config/hanzo/mcp-settings.json
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
```bash
|
||||
# Agent configuration
|
||||
export HANZO_AGENT_MODEL="openai/gpt-4"
|
||||
export HANZO_API_KEY="your-key"
|
||||
|
||||
# Tool settings
|
||||
export HANZO_COMMAND_TIMEOUT=300
|
||||
export HANZO_ENABLE_AGENT_TOOL=true
|
||||
```
|
||||
|
||||
## Project-Specific Config
|
||||
Create `.hanzo/config.json` in your project:
|
||||
```json
|
||||
{
|
||||
"enabled_tools": ["lsp", "vector_search"],
|
||||
"disabled_tools": ["sql_query"],
|
||||
"agent": {
|
||||
"model": "claude-3-sonnet",
|
||||
"max_iterations": 15
|
||||
}
|
||||
}
|
||||
```"""
|
||||
|
||||
NETWORK_AGENT_GUIDE_PROMPT = """# Network Agent Orchestration Guide
|
||||
|
||||
## Distributed AI Processing
|
||||
|
||||
The `network` tool (also accessible as `swarm` for compatibility) enables distributed AI workloads across multiple agents.
|
||||
|
||||
### Basic Multi-Agent Setup
|
||||
```python
|
||||
# Launch parallel agents
|
||||
network(
|
||||
task="Analyze this codebase for security, performance, and quality",
|
||||
agents=["security_expert", "performance_analyst", "code_reviewer"],
|
||||
mode="parallel"
|
||||
)
|
||||
```
|
||||
|
||||
### Execution Modes
|
||||
|
||||
#### 1. Local Mode (Privacy-First)
|
||||
```python
|
||||
network(
|
||||
task="Process sensitive data",
|
||||
mode="local", # Uses only local compute
|
||||
agents=["data_processor", "analyzer"]
|
||||
)
|
||||
```
|
||||
|
||||
#### 2. Distributed Mode
|
||||
```python
|
||||
network(
|
||||
task="Large-scale analysis",
|
||||
mode="distributed", # Uses network resources
|
||||
agents=["agent1", "agent2", "agent3"]
|
||||
)
|
||||
```
|
||||
|
||||
#### 3. Hybrid Mode (Default)
|
||||
```python
|
||||
network(
|
||||
task="General processing",
|
||||
mode="hybrid", # Local first, cloud fallback
|
||||
agents=["primary", "secondary"]
|
||||
)
|
||||
```
|
||||
|
||||
### Routing Strategies
|
||||
|
||||
#### Sequential Processing
|
||||
```python
|
||||
network(
|
||||
task="Multi-step workflow",
|
||||
agents=["preprocessor", "analyzer", "reporter"],
|
||||
routing="sequential" # Each agent processes in order
|
||||
)
|
||||
```
|
||||
|
||||
#### Parallel Processing
|
||||
```python
|
||||
network(
|
||||
task="Independent analyses",
|
||||
agents=["test_runner", "linter", "security_scan"],
|
||||
routing="parallel" # All agents work simultaneously
|
||||
)
|
||||
```
|
||||
|
||||
#### Consensus Decision
|
||||
```python
|
||||
network(
|
||||
task="Critical decision",
|
||||
agents=["expert1", "expert2", "expert3"],
|
||||
routing="consensus" # Agents must agree
|
||||
)
|
||||
```
|
||||
|
||||
### Claude CLI Integration
|
||||
```python
|
||||
# Run multiple Claude instances in parallel
|
||||
batch(
|
||||
description="Parallel Claude analysis",
|
||||
invocations=[
|
||||
{"tool_name": "claude_cli", "input": {"prompt": "Review architecture"}},
|
||||
{"tool_name": "claude_cli", "input": {"prompt": "Analyze performance"}},
|
||||
{"tool_name": "claude_cli", "input": {"prompt": "Check security"}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
## Advanced Features
|
||||
- **MCP Connections**: Agents communicate via MCP protocol
|
||||
- **State Sharing**: Agents share context and memory
|
||||
- **Tool Sharing**: Agents can use each other's tools
|
||||
- **Recursive Calling**: Agents can spawn sub-agents
|
||||
|
||||
## Best Practices
|
||||
1. Use local mode for sensitive data
|
||||
2. Use parallel routing for independent tasks
|
||||
3. Use consensus for critical decisions
|
||||
4. Monitor agent resource usage"""
|
||||
|
||||
PERFORMANCE_TIPS_PROMPT = """# Performance Optimization Guide
|
||||
|
||||
## Speed Optimization
|
||||
|
||||
### 1. Use Batch for Parallel Operations
|
||||
```python
|
||||
# SLOW - Sequential execution
|
||||
result1 = read(file="/file1.py")
|
||||
result2 = read(file="/file2.py")
|
||||
result3 = read(file="/file3.py")
|
||||
|
||||
# FAST - Parallel execution
|
||||
batch(
|
||||
description="Read multiple files",
|
||||
invocations=[
|
||||
{"tool_name": "read", "input": {"file_path": "/file1.py"}},
|
||||
{"tool_name": "read", "input": {"file_path": "/file2.py"}},
|
||||
{"tool_name": "read", "input": {"file_path": "/file3.py"}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### 2. Use Unified Search
|
||||
```python
|
||||
# Combines multiple search methods efficiently
|
||||
search(
|
||||
pattern="important_function",
|
||||
path="/project",
|
||||
max_results=20
|
||||
)
|
||||
```
|
||||
|
||||
### 3. Limit Result Sets
|
||||
```python
|
||||
# Use pagination
|
||||
find(pattern="*.py", page=1, page_size=50)
|
||||
|
||||
# Set max results
|
||||
grep(pattern="TODO", max_results=100)
|
||||
```
|
||||
|
||||
### 4. Use Appropriate Tools
|
||||
```python
|
||||
# For code structure - use AST search
|
||||
grep_ast(pattern="class.*Service", path="/src")
|
||||
|
||||
# For exact text - use grep
|
||||
grep(pattern="ERROR:", path="/logs")
|
||||
|
||||
# For concepts - use vector search
|
||||
vector_search(query="authentication flow", path="/src")
|
||||
```
|
||||
|
||||
## Memory Optimization
|
||||
|
||||
### 1. Process Large Files in Chunks
|
||||
```python
|
||||
# Read with offset and limit
|
||||
read(file="/large_file.log", offset=1000, limit=100)
|
||||
```
|
||||
|
||||
### 2. Use Streaming for Long Operations
|
||||
```python
|
||||
# Background long-running processes
|
||||
bash("npm run dev", background=true)
|
||||
|
||||
# Check status separately
|
||||
process(action="list")
|
||||
```
|
||||
|
||||
## Network Optimization
|
||||
|
||||
### 1. Local-First Processing
|
||||
```python
|
||||
# Use local compute when possible
|
||||
network(task="analyze", mode="local")
|
||||
```
|
||||
|
||||
### 2. Cache Results
|
||||
```python
|
||||
# Index once, search many times
|
||||
vector_index(path="/project")
|
||||
# Now searches are fast
|
||||
vector_search(query="pattern", use_cache=true)
|
||||
```
|
||||
|
||||
## Tips
|
||||
- Profile before optimizing
|
||||
- Use batch for I/O operations
|
||||
- Limit data transferred
|
||||
- Use appropriate search methods
|
||||
- Cache expensive operations"""
|
||||
|
||||
SECURITY_BEST_PRACTICES_PROMPT = """# Security Best Practices
|
||||
|
||||
## Safe Tool Usage
|
||||
|
||||
### 1. Path Validation
|
||||
```python
|
||||
# Always use absolute paths from known locations
|
||||
read(file="/home/user/project/file.py") # Good
|
||||
read(file="../../../etc/passwd") # Blocked
|
||||
|
||||
# Tools validate paths against allowed directories
|
||||
```
|
||||
|
||||
### 2. Command Injection Prevention
|
||||
```python
|
||||
# Use parameterized commands
|
||||
bash(f"grep {pattern} file.txt") # Risky if pattern is user input
|
||||
|
||||
# Better: Use dedicated tools
|
||||
grep(pattern=user_input, path="file.txt") # Safe
|
||||
```
|
||||
|
||||
### 3. Sensitive Data Handling
|
||||
```python
|
||||
# Never log sensitive data
|
||||
think("Processing user data [REDACTED]")
|
||||
|
||||
# Use local mode for sensitive operations
|
||||
network(
|
||||
task="Process PII data",
|
||||
mode="local", # Stays on device
|
||||
require_local=true
|
||||
)
|
||||
```
|
||||
|
||||
## Permission Management
|
||||
|
||||
### 1. Tool Permissions
|
||||
- Read operations are generally safe
|
||||
- Write operations require confirmation
|
||||
- System commands are restricted
|
||||
|
||||
### 2. Agent Permissions
|
||||
```python
|
||||
# Agents inherit permission restrictions
|
||||
agent(
|
||||
prompt="Analyze code",
|
||||
permissions=["read", "search"] # Limited permissions
|
||||
)
|
||||
```
|
||||
|
||||
## Data Protection
|
||||
|
||||
### 1. Local Processing
|
||||
```python
|
||||
# Keep data local
|
||||
network(mode="local", task="Process confidential data")
|
||||
```
|
||||
|
||||
### 2. Secure Communication
|
||||
- All MCP connections use secure channels
|
||||
- Agent communications are encrypted
|
||||
- No data persisted without permission
|
||||
|
||||
## Best Practices
|
||||
1. Validate all inputs
|
||||
2. Use least privilege principle
|
||||
3. Keep sensitive data local
|
||||
4. Audit tool usage
|
||||
5. Review agent outputs
|
||||
6. Don't execute untrusted code
|
||||
7. Use sandboxed environments for testing"""
|
||||
|
||||
TROUBLESHOOTING_GUIDE_PROMPT = """# Troubleshooting Guide
|
||||
|
||||
## Common Issues and Solutions
|
||||
|
||||
### 1. Tool Not Found
|
||||
```
|
||||
Error: Tool 'X' not found
|
||||
```
|
||||
**Solution:**
|
||||
```python
|
||||
# Check if tool is enabled
|
||||
tool_list()
|
||||
|
||||
# Enable if needed
|
||||
tool_enable(name="tool_name")
|
||||
```
|
||||
|
||||
### 2. Permission Denied
|
||||
```
|
||||
Error: Permission denied for path X
|
||||
```
|
||||
**Solution:**
|
||||
- Ensure path is in allowed directories
|
||||
- Launch with: `hanzo-mcp --allow-path /your/path`
|
||||
|
||||
### 3. Timeout Errors
|
||||
```
|
||||
Error: Command timed out
|
||||
```
|
||||
**Solution:**
|
||||
```python
|
||||
# Increase timeout
|
||||
bash("long_command", timeout=300)
|
||||
|
||||
# Or run in background
|
||||
bash("long_command", background=true)
|
||||
```
|
||||
|
||||
### 4. Memory Issues
|
||||
```
|
||||
Error: Result too large
|
||||
```
|
||||
**Solution:**
|
||||
```python
|
||||
# Use pagination
|
||||
find(pattern="*", page=1, page_size=50)
|
||||
|
||||
# Limit results
|
||||
grep(pattern="text", max_results=100)
|
||||
|
||||
# Read in chunks
|
||||
read(file="/large.txt", offset=0, limit=1000)
|
||||
```
|
||||
|
||||
### 5. Agent Failures
|
||||
```
|
||||
Error: Agent failed to complete task
|
||||
```
|
||||
**Solution:**
|
||||
```python
|
||||
# Increase iterations
|
||||
agent(
|
||||
prompt="task",
|
||||
max_iterations=20,
|
||||
max_tool_uses=50
|
||||
)
|
||||
|
||||
# Use simpler prompts
|
||||
# Break complex tasks into steps
|
||||
```
|
||||
|
||||
### 6. Network Issues
|
||||
```
|
||||
Error: Cannot connect to cluster
|
||||
```
|
||||
**Solution:**
|
||||
```python
|
||||
# Use local mode
|
||||
network(mode="local", task="...")
|
||||
|
||||
# Check cluster status
|
||||
mcp_stats()
|
||||
```
|
||||
|
||||
## Debugging Tips
|
||||
|
||||
### 1. Enable Verbose Output
|
||||
```python
|
||||
# Get detailed information
|
||||
stats()
|
||||
|
||||
# Check process status
|
||||
process(action="list")
|
||||
|
||||
# View logs
|
||||
process(action="logs", id="process_id")
|
||||
```
|
||||
|
||||
### 2. Test Incrementally
|
||||
```python
|
||||
# Start simple
|
||||
read(file="/test.txt")
|
||||
|
||||
# Then expand
|
||||
batch(invocations=[...])
|
||||
```
|
||||
|
||||
### 3. Check Documentation
|
||||
```python
|
||||
# View help for specific category
|
||||
"/hanzo:Filesystem tools help"
|
||||
"/hanzo:Agent tools help"
|
||||
|
||||
# Explore all tools
|
||||
"/hanzo:Explore all tools"
|
||||
```
|
||||
|
||||
## Getting Help
|
||||
1. Check error messages carefully
|
||||
2. Use `think()` to reason about issues
|
||||
3. Consult category-specific help
|
||||
4. Try simpler alternatives
|
||||
5. Report persistent issues"""
|
||||
|
||||
# Export all prompts
|
||||
__all__ = [
|
||||
"QUICK_START_PROMPT",
|
||||
"PAGINATION_GUIDE_PROMPT",
|
||||
"MEMORY_VECTOR_HELP_PROMPT",
|
||||
"DATABASE_TOOLS_HELP_PROMPT",
|
||||
"LSP_TOOLS_HELP_PROMPT",
|
||||
"CONFIGURATION_GUIDE_PROMPT",
|
||||
"NETWORK_AGENT_GUIDE_PROMPT",
|
||||
"PERFORMANCE_TIPS_PROMPT",
|
||||
"SECURITY_BEST_PRACTICES_PROMPT",
|
||||
"TROUBLESHOOTING_GUIDE_PROMPT",
|
||||
]
|
||||
@@ -1,41 +0,0 @@
|
||||
"""Example custom prompt for demonstration."""
|
||||
|
||||
# This is the actual prompt text that will be sent to Claude
|
||||
EXAMPLE_GREETING_PROMPT = """You are a friendly assistant. Please greet the user warmly and ask how you can help them today.
|
||||
Make your greeting personalized based on the time of day if possible."""
|
||||
|
||||
EXAMPLE_CODE_REVIEW_PROMPT = """Please perform a thorough code review of the most recent code changes in our conversation.
|
||||
Focus on:
|
||||
1. Code quality and best practices
|
||||
2. Potential bugs or edge cases
|
||||
3. Performance considerations
|
||||
4. Security implications
|
||||
5. Suggestions for improvement
|
||||
|
||||
Be constructive and specific in your feedback."""
|
||||
|
||||
|
||||
# You can also create dynamic prompts that take parameters
|
||||
def create_custom_analysis_prompt(file_path: str, analysis_type: str = "general"):
|
||||
"""Create a dynamic analysis prompt for a specific file."""
|
||||
|
||||
analysis_types = {
|
||||
"general": "Provide a general analysis including structure, purpose, and quality",
|
||||
"security": "Focus on security vulnerabilities and best practices",
|
||||
"performance": "Analyze performance bottlenecks and optimization opportunities",
|
||||
"refactor": "Suggest refactoring opportunities to improve code maintainability",
|
||||
}
|
||||
|
||||
analysis_instruction = analysis_types.get(analysis_type, analysis_types["general"])
|
||||
|
||||
return f"""Please analyze the file at {file_path}.
|
||||
|
||||
{analysis_instruction}
|
||||
|
||||
Your analysis should include:
|
||||
1. Overview of the file's purpose and structure
|
||||
2. Key findings based on the analysis type
|
||||
3. Specific recommendations with code examples where applicable
|
||||
4. Priority ranking of any issues found
|
||||
|
||||
Be thorough but concise in your analysis."""
|
||||
@@ -1,156 +0,0 @@
|
||||
PROJECT_SYSTEM_PROMPT = """Your are assisting me with a project.
|
||||
|
||||
Here is useful information about the environment you are running in:
|
||||
|
||||
<env>
|
||||
Working directory: {working_directory} (You need cd to this directory by yourself)
|
||||
Is directory a git repo: {is_git_repo}
|
||||
Platform: {platform}
|
||||
OS Version: {os_version}
|
||||
</env>
|
||||
|
||||
<project_info>
|
||||
directoryStructure: Below is a snapshot of this project's file structure at the start of the conversation. This snapshot will NOT update during the conversation. It skips over .gitignore patterns.
|
||||
|
||||
{directory_structure}
|
||||
|
||||
gitStatus: This is the git status at the start of the conversation. Note that this status is a snapshot in time, and will not update during the conversation.
|
||||
|
||||
Current branch: {current_branch}
|
||||
|
||||
Main branch (you will usually use this for PRs): {main_branch}
|
||||
|
||||
Status:
|
||||
|
||||
{git_status}
|
||||
|
||||
Recent commits:
|
||||
|
||||
{recent_commits}
|
||||
</project_info>
|
||||
|
||||
<available_tools>
|
||||
Hanzo AI provides 65+ tools organized by category. Key tools include:
|
||||
|
||||
# File Operations
|
||||
- read, write, edit, multi_edit: File manipulation
|
||||
- tree, find: Navigation and discovery
|
||||
|
||||
# Search & Analysis
|
||||
- grep: Fast text search
|
||||
- symbols: AST-aware symbol search
|
||||
- search: Multi-modal intelligent search
|
||||
- git_search: Git history search
|
||||
- vector_search: Semantic search
|
||||
|
||||
# Shell & Process
|
||||
- run_command: Execute commands
|
||||
- processes, pkill: Process management
|
||||
- npx, uvx: Run packages directly
|
||||
|
||||
# Development
|
||||
- jupyter: Notebook operations (read/edit actions)
|
||||
- todo: Task management (read/write actions)
|
||||
- agent: Delegate complex tasks
|
||||
- llm: Query LLMs (query/list/consensus actions)
|
||||
|
||||
# Databases
|
||||
- sql: SQL operations (query/search/stats actions)
|
||||
- graph: Graph operations (add/remove/query/search/stats actions)
|
||||
|
||||
# System
|
||||
- config: Configuration management
|
||||
- stats: Usage statistics
|
||||
- tool_enable/disable: Dynamic tool control
|
||||
|
||||
Tools follow the principle of one tool per task with multiple actions where appropriate.
|
||||
</available_tools>
|
||||
|
||||
<preferences>
|
||||
IMPORTANT: Always use the todo_write tool to plan and track tasks throughout the conversation.
|
||||
|
||||
# Code References
|
||||
When referencing specific functions or pieces of code include the pattern `file_path:line_number` to allow me to easily navigate to the source code location.
|
||||
|
||||
<example>
|
||||
user: Where are errors from the client handled?
|
||||
assistant: Clients are marked as failed in the `connectToServer` function in src/services/process.ts:712.
|
||||
</example>
|
||||
|
||||
Do what has been asked; nothing more, nothing less.
|
||||
ALWAYS prefer editing an existing file to creating a new one.
|
||||
NEVER proactively create documentation files (*.md) or README files. Only create documentation files if explicitly requested by the User.
|
||||
|
||||
# Proactiveness
|
||||
You are allowed to be proactive, but only when I ask you to do something. You should strive to strike a balance between:
|
||||
1. Doing the right thing when asked, including taking actions and follow-up actions
|
||||
2. Not surprising me with actions you take without asking
|
||||
For example, if I ask you how to approach something, you should do your best to answer their question first, and not immediately jump into taking actions.
|
||||
3. Do not add additional code explanation summary unless requested by me. After working on a file, just stop, rather than providing an explanation of what you did.
|
||||
|
||||
# Following conventions
|
||||
When making changes to files, first understand the file's code conventions. Mimic code style, use existing libraries and utilities, and follow existing patterns.
|
||||
- NEVER assume that a given library is available, even if it is well known. Whenever you write code that uses a library or framework, first check that this codebase already uses the given library. For example, you might look at neighboring files, or check the package.json (or cargo.toml, and so on depending on the language).
|
||||
- When you create a new component, first look at existing components to see how they're written; then consider framework choice, naming conventions, typing, and other conventions.
|
||||
- When you edit a piece of code, first look at the code's surrounding context (especially its imports) to understand the code's choice of frameworks and libraries. Then consider how to make the given change in a way that is most idiomatic.
|
||||
- Always follow security best practices. Never introduce code that exposes or logs secrets and keys. Never commit secrets or keys to the repository.
|
||||
</preferences>
|
||||
|
||||
|
||||
<task_management>
|
||||
# Task Management
|
||||
You have access to the todo tool (actions: read, write) to help you manage and plan tasks. Use this tool VERY frequently to ensure that you are tracking your tasks and giving me visibility into your progress.
|
||||
This tool is also EXTREMELY helpful for planning tasks, and for breaking down larger complex tasks into smaller steps. If you do not use this tool when planning, you may forget to do important tasks - and that is unacceptable.
|
||||
|
||||
It is critical that you mark todos as completed as soon as you are done with a task. Do not batch up multiple tasks before marking them as completed.
|
||||
|
||||
Examples:
|
||||
<example>
|
||||
user: Run the build and fix any type errors
|
||||
assistant: I'm going to use the todo tool with write action to add the following items to the todo list:
|
||||
- Run the build
|
||||
- Fix any type errors
|
||||
|
||||
I'm now going to run the build using Bash.
|
||||
|
||||
Looks like I found 10 type errors. I'm going to use the todo tool with write action to add 10 items to the todo list.
|
||||
|
||||
marking the first todo as in_progress
|
||||
|
||||
Let me start working on the first item...
|
||||
|
||||
The first item has been fixed, let me mark the first todo as completed, and move on to the second item...
|
||||
</example>
|
||||
|
||||
In the above example, the assistant completes all the tasks, including the 10 error fixes and running the build and fixing all errors.
|
||||
|
||||
<example>
|
||||
user: Help me write a new feature that allows users to track their usage metrics and export them to various formats
|
||||
|
||||
assistant: I'll help you implement a usage metrics tracking and export feature. Let me first use the todo tool with write action to plan this task.
|
||||
Adding the following todos to the todo list:
|
||||
1. Research existing metrics tracking in the codebase
|
||||
2. Design the metrics collection system
|
||||
3. Implement core metrics tracking functionality
|
||||
4. Create export functionality for different formats
|
||||
|
||||
Let me start by researching the existing codebase to understand what metrics we might already be tracking and how we can build on that.
|
||||
|
||||
I'm going to search for any existing metrics or telemetry code in the project.
|
||||
|
||||
I've found some existing telemetry code. Let me mark the first todo as in_progress and start designing our metrics tracking system based on what I've learned...
|
||||
|
||||
[Assistant continues implementing the feature step by step, marking todos as in_progress and completed as they go]
|
||||
</example>
|
||||
|
||||
# Doing tasks
|
||||
I will primarily request you perform software engineering tasks. This includes solving bugs, adding new functionality, refactoring code, explaining code, and more. For these tasks the following steps are recommended:
|
||||
- Use the todo tool with write action to plan the task if required
|
||||
- Use the available search tools (grep, symbols, search, git_search) to understand the codebase and my query. The 'search' tool intelligently combines multiple search strategies. You are encouraged to use search tools extensively both in parallel and sequentially.
|
||||
- Implement the solution using all tools available to you
|
||||
- Verify the solution if possible with tests. NEVER assume specific test framework or test script. Check the README or search codebase to determine the testing approach.
|
||||
- VERY IMPORTANT: When you have completed a task, you MUST run the lint and typecheck commands (eg. npm run lint, npm run typecheck, ruff, etc.) with Bash if they were provided to you to ensure your code is correct. If you are unable to find the correct command, ask me for the command to run and if they supply it, proactively suggest writing it to CLAUDE.md so that you will know to run it next time.
|
||||
NEVER commit changes unless I explicitly ask you to. It is VERY IMPORTANT to only commit when explicitly asked, otherwise I will feel that you are being too proactive.
|
||||
|
||||
- Tool results and user messages may include <system-reminder> tags. <system-reminder> tags contain useful information and reminders.
|
||||
</task_management>"""
|
||||
@@ -1,110 +0,0 @@
|
||||
from typing import Any
|
||||
|
||||
# Import TodoStorage to access todo data
|
||||
from hanzo_mcp.tools.todo.base import TodoStorage
|
||||
|
||||
PROJECT_TODO_EMPTY_REMINDER = """<system-reminder>This is a reminder that your todo list is currently empty. DO NOT mention this to me explicitly because i have already aware. If you are working on tasks that would benefit from a todo list please use the todo_write tool to create one. If not, please feel free to ignore.</system-reminder>"""
|
||||
|
||||
|
||||
PROJECT_TODO_REMINDER = """<system-reminder>
|
||||
This is a reminder that you have a to-do list for this project. The to-do list session ID is: {session_id}. You can use the todo_write tool to add new to-dos to the list.
|
||||
|
||||
The to-do list is shown below, so you do not need to read it using the todo_read tool before your next time using the todo_write tool:
|
||||
|
||||
{todo_list}
|
||||
|
||||
</system-reminder>"""
|
||||
|
||||
|
||||
def format_todo_list_concise(todos: list[dict[str, Any]]) -> str:
|
||||
"""Format a todo list in a concise format for inclusion in prompts.
|
||||
|
||||
Args:
|
||||
todos: List of todo items
|
||||
|
||||
Returns:
|
||||
Formatted string representation of the todo list
|
||||
"""
|
||||
if not todos:
|
||||
return "No todos found."
|
||||
|
||||
formatted_lines = []
|
||||
for todo in todos:
|
||||
status = todo.get("status", "unknown")
|
||||
priority = todo.get("priority", "medium")
|
||||
content = todo.get("content", "No content")
|
||||
todo_id = todo.get("id", "no-id")
|
||||
|
||||
# Handle empty strings as well as missing values
|
||||
if not content or not str(content).strip():
|
||||
content = "No content"
|
||||
if not todo_id or not str(todo_id).strip():
|
||||
todo_id = "no-id"
|
||||
|
||||
# Create status indicator
|
||||
status_indicator = {
|
||||
"pending": "[ ]",
|
||||
"in_progress": "[~]",
|
||||
"completed": "[✓]",
|
||||
}.get(status, "[?]")
|
||||
|
||||
# Create priority indicator
|
||||
priority_indicator = {"high": "🔴", "medium": "🟡", "low": "🟢"}.get(
|
||||
priority, "⚪"
|
||||
)
|
||||
|
||||
formatted_lines.append(
|
||||
f"{status_indicator} {priority_indicator} {content} (id: {todo_id})"
|
||||
)
|
||||
|
||||
return "\n".join(formatted_lines)
|
||||
|
||||
|
||||
def has_unfinished_todos(todos: list[dict[str, Any]]) -> bool:
|
||||
"""Check if there are any unfinished todos in the list.
|
||||
|
||||
Args:
|
||||
todos: List of todo items
|
||||
|
||||
Returns:
|
||||
True if there are unfinished todos, False otherwise
|
||||
"""
|
||||
if not todos:
|
||||
return False
|
||||
|
||||
for todo in todos:
|
||||
status = todo.get("status", "pending")
|
||||
if status in ["pending", "in_progress"]:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def get_project_todo_reminder(session_id: str | None = None) -> str:
|
||||
"""Get the appropriate todo reminder for a session.
|
||||
|
||||
Args:
|
||||
session_id: Session ID to check todos for. If None, finds the latest active session.
|
||||
|
||||
Returns:
|
||||
Either PROJECT_TODO_EMPTY_REMINDER or PROJECT_TODO_REMINDER with formatted content
|
||||
"""
|
||||
# If no session_id provided, try to find the latest active session
|
||||
if session_id is None:
|
||||
session_id = TodoStorage.find_latest_active_session()
|
||||
if session_id is None:
|
||||
# No active sessions found
|
||||
return PROJECT_TODO_EMPTY_REMINDER
|
||||
|
||||
# Get todos for the session
|
||||
todos = TodoStorage.get_todos(session_id)
|
||||
|
||||
# Check if we have unfinished todos
|
||||
if not has_unfinished_todos(todos):
|
||||
return PROJECT_TODO_EMPTY_REMINDER
|
||||
|
||||
# Format the todo list and return the reminder with content
|
||||
formatted_todos = format_todo_list_concise(todos)
|
||||
return PROJECT_TODO_REMINDER.format(
|
||||
session_id=session_id, todo_list=formatted_todos
|
||||
)
|
||||
@@ -1,606 +0,0 @@
|
||||
"""Tool explorer prompts for discovering and using Hanzo MCP tools."""
|
||||
|
||||
TOOL_EXPLORER_PROMPT = """# Hanzo MCP Tool Explorer
|
||||
|
||||
You have access to a comprehensive suite of tools through the Hanzo MCP system. These tools can be used individually or combined using the batch tool for powerful multi-agent workflows.
|
||||
|
||||
## Tool Categories
|
||||
|
||||
### 🤖 Agent Tools
|
||||
Tools for delegating tasks to specialized AI agents:
|
||||
- **dispatch_agent**: Launch a specialized agent for file exploration and analysis
|
||||
- **swarm**: Orchestrate multiple agents working together
|
||||
- **claude_cli**: Interact with Claude via CLI
|
||||
- **critic**: Critical analysis and review
|
||||
- **review**: Code review and feedback
|
||||
- **iching**: Decision-making guidance using I Ching
|
||||
|
||||
### 📁 Filesystem Tools
|
||||
Tools for file and directory operations:
|
||||
- **read_files**: Read one or multiple files
|
||||
- **write_file**: Create or overwrite files
|
||||
- **edit_file**: Make precise edits to files
|
||||
- **multi_edit**: Multiple edits in one operation
|
||||
- **directory_tree**: View directory structure
|
||||
- **find**: Find files by pattern
|
||||
- **grep**: Search file contents
|
||||
- **grep_ast**: Search with AST context
|
||||
- **search_content**: Unified search across multiple methods
|
||||
- **symbols**: Find code symbols
|
||||
- **content_replace**: Replace patterns across files
|
||||
- **batch_search**: Run multiple searches in parallel
|
||||
|
||||
### 🐚 Shell Tools
|
||||
Tools for command execution:
|
||||
- **run_command**: Execute shell commands
|
||||
- **bash**: Run bash commands with session persistence
|
||||
- **npx**: Run Node packages
|
||||
- **uvx**: Run Python packages
|
||||
- **process**: Manage background processes
|
||||
- **open**: Open files/URLs in default app
|
||||
|
||||
### 🧠 AI/LLM Tools
|
||||
Tools for AI operations:
|
||||
- **llm**: Query various LLM providers
|
||||
- **consensus**: Get consensus from multiple models
|
||||
- **think**: Structured thinking and planning
|
||||
- **critic**: Critical analysis
|
||||
|
||||
### 💾 Database Tools
|
||||
Tools for data operations:
|
||||
- **sql_query**: Execute SQL queries
|
||||
- **graph_add**: Add to graph database
|
||||
- **vector_search**: Semantic search
|
||||
- **index**: Manage search indices
|
||||
|
||||
### 📓 Jupyter Tools
|
||||
Tools for notebook operations:
|
||||
- **notebook_read**: Read Jupyter notebooks
|
||||
- **notebook_edit**: Edit notebook cells
|
||||
|
||||
### ✅ Todo Tools
|
||||
Tools for task management:
|
||||
- **todo**: Manage todo lists
|
||||
- **todo_read**: Read current todos
|
||||
- **todo_write**: Update todo items
|
||||
|
||||
### 🔧 Configuration Tools
|
||||
Tools for settings and configuration:
|
||||
- **config**: Manage tool configuration
|
||||
- **mode**: Switch developer modes
|
||||
- **tool_list**: List available tools
|
||||
- **tool_enable/disable**: Toggle tools
|
||||
|
||||
### 🔍 LSP Tools
|
||||
Language Server Protocol tools:
|
||||
- **lsp**: Code intelligence operations
|
||||
|
||||
### 🌐 MCP Tools
|
||||
Model Context Protocol management:
|
||||
- **mcp_add**: Add MCP servers
|
||||
- **mcp_remove**: Remove MCP servers
|
||||
- **mcp_stats**: View MCP statistics
|
||||
|
||||
## Using Tools with Batch
|
||||
|
||||
The batch tool allows you to run multiple tools in parallel for maximum efficiency:
|
||||
|
||||
```python
|
||||
batch(
|
||||
description="Analyze project structure",
|
||||
invocations=[
|
||||
{"tool_name": "directory_tree", "input": {"path": "/project"}},
|
||||
{"tool_name": "grep", "input": {"pattern": "TODO", "path": "/project"}},
|
||||
{"tool_name": "find", "input": {"pattern": "*.test.js", "path": "/project"}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
## Tool Usage Examples
|
||||
|
||||
### Example 1: Code Analysis Workflow
|
||||
```python
|
||||
# First, explore the project structure
|
||||
directory_tree(path="/project", depth=3)
|
||||
|
||||
# Search for specific patterns
|
||||
batch(
|
||||
description="Find all API endpoints",
|
||||
invocations=[
|
||||
{"tool_name": "grep", "input": {"pattern": "app\\.(get|post|put|delete)", "path": "/project/src"}},
|
||||
{"tool_name": "grep_ast", "input": {"pattern": "router", "path": "/project/src"}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### Example 2: Multi-Agent Analysis
|
||||
```python
|
||||
# Dispatch specialized agents for different tasks
|
||||
batch(
|
||||
description="Comprehensive code analysis",
|
||||
invocations=[
|
||||
{"tool_name": "dispatch_agent", "input": {"prompt": "Analyze security vulnerabilities in /project/src"}},
|
||||
{"tool_name": "dispatch_agent", "input": {"prompt": "Find performance bottlenecks in database queries"}},
|
||||
{"tool_name": "dispatch_agent", "input": {"prompt": "Review test coverage and suggest improvements"}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### Example 3: Refactoring Workflow
|
||||
```python
|
||||
# Find all instances of a pattern
|
||||
search_content(pattern="oldFunction", path="/project")
|
||||
|
||||
# Review the code
|
||||
critic(analysis="Review the usage of oldFunction and suggest refactoring approach")
|
||||
|
||||
# Make the changes
|
||||
batch(
|
||||
description="Refactor oldFunction to newFunction",
|
||||
invocations=[
|
||||
{"tool_name": "content_replace", "input": {
|
||||
"pattern": "oldFunction",
|
||||
"replacement": "newFunction",
|
||||
"path": "/project/src"
|
||||
}},
|
||||
{"tool_name": "run_command", "input": {"command": "npm test"}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
## Best Practices
|
||||
|
||||
1. **Use batch for parallel operations**: When you need to run multiple independent operations, use batch to run them concurrently.
|
||||
|
||||
2. **Combine search tools**: Use unified_search for comprehensive results, grep for simple text matching, and grep_ast for code structure understanding.
|
||||
|
||||
3. **Leverage agents for complex tasks**: Use dispatch_agent for tasks requiring deep analysis or multiple steps.
|
||||
|
||||
4. **Track progress with todos**: Use todo tools to manage multi-step workflows.
|
||||
|
||||
5. **Think before acting**: Use the think tool to plan complex operations.
|
||||
|
||||
## Getting Started
|
||||
|
||||
To explore available tools in detail:
|
||||
1. Use `tool_list()` to see all available tools
|
||||
2. Each tool has detailed documentation in its implementation
|
||||
3. Tools can be combined creatively for powerful workflows
|
||||
|
||||
Would you like to explore any specific tool category or see more examples?"""
|
||||
|
||||
# Tool category specific prompts
|
||||
FILESYSTEM_TOOLS_HELP = """# Filesystem Tools Guide
|
||||
|
||||
## Core File Operations
|
||||
|
||||
### Reading Files
|
||||
```python
|
||||
# Read a single file
|
||||
read_files(paths=["/path/to/file.py"])
|
||||
|
||||
# Read multiple files at once
|
||||
read_files(paths=[
|
||||
"/project/src/main.py",
|
||||
"/project/src/utils.py",
|
||||
"/project/tests/test_main.py"
|
||||
])
|
||||
|
||||
# Read with line limits
|
||||
read_files(paths=["/large/file.py"], lines=100, offset=500)
|
||||
```
|
||||
|
||||
### Editing Files
|
||||
```python
|
||||
# Simple edit
|
||||
edit_file(
|
||||
path="/src/main.py",
|
||||
edits=[{"oldText": "old code", "newText": "new code"}]
|
||||
)
|
||||
|
||||
# Multiple edits in one file
|
||||
multi_edit(
|
||||
file_path="/src/utils.py",
|
||||
edits=[
|
||||
{"old_string": "import old", "new_string": "import new"},
|
||||
{"old_string": "old_function", "new_string": "new_function"}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### Searching
|
||||
```python
|
||||
# Find files by pattern
|
||||
find(pattern="*.test.js", path="/project")
|
||||
|
||||
# Search file contents
|
||||
grep(pattern="TODO|FIXME", path="/project", include="*.py")
|
||||
|
||||
# Search with AST context
|
||||
grep_ast(pattern="class.*Controller", path="/project/src")
|
||||
|
||||
# Unified search (combines multiple search methods)
|
||||
search_content(
|
||||
pattern="authentication",
|
||||
path="/project",
|
||||
enable_vector=True,
|
||||
enable_ast=True
|
||||
)
|
||||
```
|
||||
|
||||
### Batch Operations
|
||||
```python
|
||||
# Search across multiple patterns simultaneously
|
||||
batch_search(
|
||||
queries=[
|
||||
{"pattern": "login", "type": "text"},
|
||||
{"pattern": "authenticate", "type": "semantic"},
|
||||
{"pattern": "class.*Auth", "type": "ast"}
|
||||
],
|
||||
path="/project"
|
||||
)
|
||||
```
|
||||
|
||||
## Advanced Features
|
||||
|
||||
### Content Replacement
|
||||
```python
|
||||
# Replace across multiple files
|
||||
content_replace(
|
||||
pattern="oldAPI",
|
||||
replacement="newAPI",
|
||||
path="/project/src",
|
||||
dry_run=True # Preview changes first
|
||||
)
|
||||
```
|
||||
|
||||
### Directory Exploration
|
||||
```python
|
||||
# View directory structure
|
||||
directory_tree(path="/project", depth=3, include_filtered=False)
|
||||
|
||||
# Find specific file types
|
||||
find(
|
||||
pattern="*.py",
|
||||
path="/project",
|
||||
min_size="1KB",
|
||||
max_size="100KB",
|
||||
modified_after="1 week ago"
|
||||
)
|
||||
```"""
|
||||
|
||||
AGENT_TOOLS_HELP = """# Agent Tools Guide
|
||||
|
||||
## Dispatching Agents
|
||||
|
||||
The agent tools allow you to delegate complex tasks to specialized sub-agents that have access to file operations and search capabilities.
|
||||
|
||||
### Basic Agent Dispatch
|
||||
```python
|
||||
# Dispatch a single agent for analysis
|
||||
dispatch_agent(
|
||||
prompt="Analyze the authentication system in /project/src/auth and identify security vulnerabilities"
|
||||
)
|
||||
|
||||
# Multiple agents for different aspects
|
||||
batch(
|
||||
description="Comprehensive security audit",
|
||||
invocations=[
|
||||
{"tool_name": "dispatch_agent", "input": {
|
||||
"prompt": "Review authentication implementation for security issues"
|
||||
}},
|
||||
{"tool_name": "dispatch_agent", "input": {
|
||||
"prompt": "Check for SQL injection vulnerabilities in database queries"
|
||||
}},
|
||||
{"tool_name": "dispatch_agent", "input": {
|
||||
"prompt": "Analyze API endpoints for authorization bypass risks"
|
||||
}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### Swarm Operations
|
||||
```python
|
||||
# Create a swarm of agents working together
|
||||
swarm(
|
||||
agents=[
|
||||
{
|
||||
"id": "analyzer",
|
||||
"role": "Code Analyzer",
|
||||
"goal": "Identify code quality issues",
|
||||
"backstory": "Expert in clean code principles",
|
||||
"tools": ["grep_ast", "read_files", "symbols"]
|
||||
},
|
||||
{
|
||||
"id": "refactorer",
|
||||
"role": "Code Refactorer",
|
||||
"goal": "Suggest and implement improvements",
|
||||
"backstory": "Specialist in code optimization",
|
||||
"tools": ["edit_file", "multi_edit", "content_replace"]
|
||||
}
|
||||
],
|
||||
tasks=[
|
||||
{
|
||||
"description": "Find code smells and anti-patterns",
|
||||
"assigned_to": "analyzer"
|
||||
},
|
||||
{
|
||||
"description": "Refactor identified issues",
|
||||
"assigned_to": "refactorer",
|
||||
"depends_on": ["analyzer"]
|
||||
}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### Specialized Agents
|
||||
|
||||
#### Critic Agent
|
||||
```python
|
||||
# Get critical analysis
|
||||
critic(
|
||||
analysis="Review this implementation for potential issues:\\n" + code_snippet
|
||||
)
|
||||
```
|
||||
|
||||
#### Review Agent
|
||||
```python
|
||||
# Comprehensive code review
|
||||
review(
|
||||
files=["/src/main.py", "/src/utils.py"],
|
||||
focus_areas=["security", "performance", "maintainability"]
|
||||
)
|
||||
```
|
||||
|
||||
#### I Ching Guidance
|
||||
```python
|
||||
# Get decision-making guidance
|
||||
iching(
|
||||
question="Should we refactor the authentication system now or after the release?"
|
||||
)
|
||||
```
|
||||
|
||||
## Agent Capabilities
|
||||
|
||||
Agents dispatched through these tools have access to:
|
||||
- File reading and searching
|
||||
- Pattern matching and AST analysis
|
||||
- Directory exploration
|
||||
- Comprehensive search capabilities
|
||||
|
||||
They cannot:
|
||||
- Modify files directly
|
||||
- Execute shell commands
|
||||
- Make external API calls
|
||||
|
||||
This makes them safe for exploration and analysis tasks."""
|
||||
|
||||
SHELL_TOOLS_HELP = """# Shell Tools Guide
|
||||
|
||||
## Command Execution
|
||||
|
||||
### Basic Commands
|
||||
```python
|
||||
# Run simple commands
|
||||
run_command(command="ls -la", cwd="/project")
|
||||
run_command(command="git status")
|
||||
|
||||
# Run with environment variables
|
||||
run_command(
|
||||
command="npm test",
|
||||
env={"NODE_ENV": "test", "CI": "true"}
|
||||
)
|
||||
```
|
||||
|
||||
### Bash Sessions
|
||||
```python
|
||||
# Use bash for session persistence
|
||||
bash(command="cd /project && npm install")
|
||||
bash(command="export API_KEY=test123 && npm run dev")
|
||||
```
|
||||
|
||||
### Package Runners
|
||||
```python
|
||||
# Run Node packages
|
||||
npx(package="prettier", args="--write src/**/*.js")
|
||||
npx(package="create-react-app", args="my-app")
|
||||
|
||||
# Run Python packages
|
||||
uvx(package="ruff", args="check .")
|
||||
uvx(package="black", args="--check src/")
|
||||
```
|
||||
|
||||
### Background Processes
|
||||
```python
|
||||
# Start long-running processes
|
||||
run_command(command="npm run dev", background=True)
|
||||
|
||||
# Manage processes
|
||||
process(action="list") # List all background processes
|
||||
process(action="logs", id="npm_abc123") # View logs
|
||||
process(action="kill", id="npm_abc123") # Stop process
|
||||
```
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
### Batch Operations
|
||||
```python
|
||||
# Run multiple commands efficiently
|
||||
batch(
|
||||
description="Run tests and linting",
|
||||
invocations=[
|
||||
{"tool_name": "run_command", "input": {"command": "npm test"}},
|
||||
{"tool_name": "run_command", "input": {"command": "npm run lint"}},
|
||||
{"tool_name": "run_command", "input": {"command": "npm audit"}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### Working with Output
|
||||
```python
|
||||
# Capture and process output
|
||||
result = run_command(command="git log --oneline -10")
|
||||
# Process result.output for analysis
|
||||
```
|
||||
|
||||
### Platform-Specific Commands
|
||||
```python
|
||||
# Open files/URLs in default application
|
||||
open(path="https://github.com/user/repo")
|
||||
open(path="/path/to/document.pdf")
|
||||
```"""
|
||||
|
||||
BATCH_TOOL_EXAMPLES = """# Batch Tool Mastery
|
||||
|
||||
The batch tool is one of the most powerful features in Hanzo MCP, allowing parallel execution of multiple tools for maximum efficiency.
|
||||
|
||||
## Basic Batch Usage
|
||||
|
||||
```python
|
||||
batch(
|
||||
description="Project setup",
|
||||
invocations=[
|
||||
{"tool_name": "run_command", "input": {"command": "npm install"}},
|
||||
{"tool_name": "run_command", "input": {"command": "pip install -r requirements.txt"}},
|
||||
{"tool_name": "directory_tree", "input": {"path": ".", "depth": 2}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
## Advanced Patterns
|
||||
|
||||
### 1. Parallel Search Operations
|
||||
```python
|
||||
batch(
|
||||
description="Find all authentication code",
|
||||
invocations=[
|
||||
{"tool_name": "grep", "input": {
|
||||
"pattern": "login|auth|session",
|
||||
"path": "/src"
|
||||
}},
|
||||
{"tool_name": "grep_ast", "input": {
|
||||
"pattern": "class.*Auth",
|
||||
"path": "/src"
|
||||
}},
|
||||
{"tool_name": "find", "input": {
|
||||
"pattern": "*auth*.py",
|
||||
"path": "/src"
|
||||
}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### 2. Multi-Agent Analysis
|
||||
```python
|
||||
batch(
|
||||
description="Comprehensive code analysis",
|
||||
invocations=[
|
||||
{"tool_name": "dispatch_agent", "input": {
|
||||
"prompt": "Analyze code quality in /src/core modules"
|
||||
}},
|
||||
{"tool_name": "dispatch_agent", "input": {
|
||||
"prompt": "Review test coverage for /src/core modules"
|
||||
}},
|
||||
{"tool_name": "dispatch_agent", "input": {
|
||||
"prompt": "Identify performance bottlenecks in database operations"
|
||||
}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### 3. File Operations
|
||||
```python
|
||||
batch(
|
||||
description="Read configuration files",
|
||||
invocations=[
|
||||
{"tool_name": "read_files", "input": {
|
||||
"paths": ["package.json", "tsconfig.json", ".env.example"]
|
||||
}},
|
||||
{"tool_name": "read_files", "input": {
|
||||
"paths": ["src/config/database.js", "src/config/auth.js"]
|
||||
}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### 4. Complex Workflows
|
||||
```python
|
||||
# Step 1: Analyze
|
||||
analysis_batch = batch(
|
||||
description="Analyze codebase",
|
||||
invocations=[
|
||||
{"tool_name": "grep", "input": {"pattern": "TODO|FIXME", "path": "."}},
|
||||
{"tool_name": "dispatch_agent", "input": {
|
||||
"prompt": "Find unused imports and dead code"
|
||||
}}
|
||||
]
|
||||
)
|
||||
|
||||
# Step 2: Based on analysis, perform fixes
|
||||
fix_batch = batch(
|
||||
description="Fix identified issues",
|
||||
invocations=[
|
||||
{"tool_name": "content_replace", "input": {
|
||||
"pattern": "old_import",
|
||||
"replacement": "new_import",
|
||||
"path": "/src"
|
||||
}},
|
||||
{"tool_name": "run_command", "input": {"command": "npm run lint:fix"}}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
## Batch Tool Best Practices
|
||||
|
||||
1. **Group Related Operations**: Batch operations that are logically related
|
||||
2. **Maximize Parallelism**: Independent operations should be in the same batch
|
||||
3. **Use Descriptive Names**: The description helps track what the batch does
|
||||
4. **Handle Results**: Each tool result is returned in order
|
||||
|
||||
## Limitations
|
||||
|
||||
- Tools in a batch cannot depend on each other's results
|
||||
- All tools run in parallel when possible
|
||||
- Maximum efficiency with truly independent operations
|
||||
|
||||
## Available Tools for Batch
|
||||
|
||||
The following tools can be used in batch operations:
|
||||
- dispatch_agent
|
||||
- read_files
|
||||
- directory_tree
|
||||
- grep
|
||||
- grep_ast
|
||||
- run_command
|
||||
- notebook_read
|
||||
- find
|
||||
- search_content
|
||||
- symbols
|
||||
- git_search
|
||||
|
||||
Tools NOT available in batch (require state/session):
|
||||
- write_file
|
||||
- edit_file
|
||||
- multi_edit
|
||||
- think
|
||||
- todo_write"""
|
||||
|
||||
|
||||
def create_tool_category_prompt(category: str, tools: list[str]):
|
||||
"""Create a dynamic prompt for a specific tool category."""
|
||||
|
||||
tool_descriptions = {
|
||||
"filesystem": FILESYSTEM_TOOLS_HELP,
|
||||
"agent": AGENT_TOOLS_HELP,
|
||||
"shell": SHELL_TOOLS_HELP,
|
||||
"batch": BATCH_TOOL_EXAMPLES,
|
||||
}
|
||||
|
||||
base_prompt = tool_descriptions.get(
|
||||
category, f"# {category.title()} Tools\n\nAvailable tools in this category:\n"
|
||||
)
|
||||
|
||||
if category not in tool_descriptions:
|
||||
base_prompt += "\n".join(f"- **{tool}**: [Tool description]" for tool in tools)
|
||||
|
||||
return base_prompt
|
||||
@@ -1,286 +0,0 @@
|
||||
import platform
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from git import Repo
|
||||
|
||||
GIT_AVAILABLE = True
|
||||
except ImportError:
|
||||
GIT_AVAILABLE = False
|
||||
|
||||
|
||||
def get_os_info() -> tuple[str, str, str]:
|
||||
"""Get the operating system information.
|
||||
Returns:
|
||||
tuple: A tuple containing the system name, release, and version.
|
||||
"""
|
||||
system = platform.system() # noqa: F821
|
||||
release = platform.release()
|
||||
version = platform.version()
|
||||
|
||||
if system == "Darwin":
|
||||
system = "MacOS"
|
||||
elif system == "Linux":
|
||||
try:
|
||||
with open("/etc/os-release") as f:
|
||||
for line in f:
|
||||
if line.startswith("NAME="):
|
||||
name = line.split("=")[1].strip().strip('"')
|
||||
if "Ubuntu" in name:
|
||||
system = "Ubuntu"
|
||||
elif "Debian" in name:
|
||||
system = "Debian"
|
||||
elif "Fedora" in name:
|
||||
system = "Fedora"
|
||||
elif "CentOS" in name:
|
||||
system = "CentOS"
|
||||
elif "Arch Linux" in name:
|
||||
system = "Arch Linux"
|
||||
system = name
|
||||
except FileNotFoundError:
|
||||
dist = platform.freedesktop_os_release()
|
||||
if dist and "NAME" in dist:
|
||||
name = dist["NAME"]
|
||||
if "Ubuntu" in name:
|
||||
system = "Ubuntu"
|
||||
system = name
|
||||
system = "Linux"
|
||||
elif system == "Java":
|
||||
system = "Java"
|
||||
|
||||
return system, release, version
|
||||
|
||||
|
||||
def get_directory_structure(
|
||||
path: str, max_depth: int = 3, include_filtered: bool = False
|
||||
) -> str:
|
||||
"""Get a directory structure similar to directory_tree tool.
|
||||
|
||||
Args:
|
||||
path: The directory path to scan
|
||||
max_depth: Maximum depth to traverse (0 for unlimited)
|
||||
include_filtered: Whether to include normally filtered directories
|
||||
|
||||
Returns:
|
||||
Formatted directory structure as a string
|
||||
"""
|
||||
try:
|
||||
dir_path = Path(path)
|
||||
|
||||
if not dir_path.exists() or not dir_path.is_dir():
|
||||
return f"Error: {path} is not a valid directory"
|
||||
|
||||
# Define filtered directories (same as directory_tree.py)
|
||||
FILTERED_DIRECTORIES = {
|
||||
".git",
|
||||
"node_modules",
|
||||
".venv",
|
||||
"venv",
|
||||
"__pycache__",
|
||||
".pytest_cache",
|
||||
".idea",
|
||||
".vs",
|
||||
".vscode",
|
||||
"dist",
|
||||
"build",
|
||||
"target",
|
||||
".ruff_cache",
|
||||
".llm-context",
|
||||
}
|
||||
|
||||
def should_filter(current_path: Path) -> bool:
|
||||
"""Check if a directory should be filtered."""
|
||||
# Don't filter if it's the explicitly requested path
|
||||
if str(current_path.absolute()) == str(dir_path.absolute()):
|
||||
return False
|
||||
# Filter based on directory name if filtering is enabled
|
||||
return current_path.name in FILTERED_DIRECTORIES and not include_filtered
|
||||
|
||||
def build_tree(current_path: Path, current_depth: int = 0) -> list[dict]:
|
||||
"""Build directory tree recursively."""
|
||||
result = []
|
||||
|
||||
try:
|
||||
# Sort entries: directories first, then files alphabetically
|
||||
entries = sorted(
|
||||
current_path.iterdir(), key=lambda x: (not x.is_dir(), x.name)
|
||||
)
|
||||
|
||||
for entry in entries:
|
||||
if entry.is_dir():
|
||||
entry_data = {"name": entry.name, "type": "directory"}
|
||||
|
||||
# Check if we should filter this directory
|
||||
if should_filter(entry):
|
||||
entry_data["skipped"] = "filtered-directory"
|
||||
result.append(entry_data)
|
||||
continue
|
||||
|
||||
# Check depth limit (if enabled)
|
||||
if max_depth > 0 and current_depth >= max_depth:
|
||||
entry_data["skipped"] = "depth-limit"
|
||||
result.append(entry_data)
|
||||
continue
|
||||
|
||||
# Process children recursively
|
||||
entry_data["children"] = build_tree(entry, current_depth + 1)
|
||||
result.append(entry_data)
|
||||
else:
|
||||
# Add files only if within depth limit
|
||||
if max_depth <= 0 or current_depth < max_depth:
|
||||
result.append({"name": entry.name, "type": "file"})
|
||||
|
||||
except Exception:
|
||||
# Skip directories we can't read
|
||||
pass
|
||||
|
||||
return result
|
||||
|
||||
def format_tree(tree_data: list[dict], level: int = 0) -> list[str]:
|
||||
"""Format tree data as indented strings."""
|
||||
lines = []
|
||||
|
||||
for item in tree_data:
|
||||
# Indentation based on level
|
||||
indent = " " * level
|
||||
|
||||
# Format based on type
|
||||
if item["type"] == "directory":
|
||||
if "skipped" in item:
|
||||
lines.append(
|
||||
f"{indent}{item['name']}/ [skipped - {item['skipped']}]"
|
||||
)
|
||||
else:
|
||||
lines.append(f"{indent}{item['name']}/")
|
||||
# Add children with increased indentation if present
|
||||
if "children" in item:
|
||||
lines.extend(format_tree(item["children"], level + 1))
|
||||
else:
|
||||
# File
|
||||
lines.append(f"{indent}{item['name']}")
|
||||
|
||||
return lines
|
||||
|
||||
# Build and format the tree
|
||||
tree_data = build_tree(dir_path)
|
||||
formatted_lines = format_tree(tree_data)
|
||||
|
||||
# Add the root directory path as a prefix
|
||||
result = f"- {dir_path}/"
|
||||
if formatted_lines:
|
||||
result += "\n" + "\n".join(f" {line}" for line in formatted_lines)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
return f"Error generating directory structure: {str(e)}"
|
||||
|
||||
|
||||
def get_git_info(path: str) -> dict[str, str | None]:
|
||||
"""Get git information for a repository.
|
||||
|
||||
Args:
|
||||
path: Path to the git repository
|
||||
|
||||
Returns:
|
||||
Dictionary containing git information
|
||||
"""
|
||||
if not GIT_AVAILABLE:
|
||||
return {
|
||||
"current_branch": None,
|
||||
"main_branch": None,
|
||||
"git_status": "GitPython not available",
|
||||
"recent_commits": "GitPython not available",
|
||||
}
|
||||
|
||||
try:
|
||||
repo = Repo(path)
|
||||
|
||||
# Get current branch
|
||||
try:
|
||||
current_branch = repo.active_branch.name
|
||||
except Exception:
|
||||
current_branch = "HEAD (detached)"
|
||||
|
||||
# Try to determine main branch
|
||||
main_branch = "main" # default
|
||||
try:
|
||||
# Check if 'main' exists
|
||||
if "origin/main" in [ref.name for ref in repo.refs]:
|
||||
main_branch = "main"
|
||||
elif "origin/master" in [ref.name for ref in repo.refs]:
|
||||
main_branch = "master"
|
||||
elif "main" in [ref.name for ref in repo.refs]:
|
||||
main_branch = "main"
|
||||
elif "master" in [ref.name for ref in repo.refs]:
|
||||
main_branch = "master"
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Get git status
|
||||
try:
|
||||
status_lines = []
|
||||
|
||||
# Check for staged changes
|
||||
staged_files = list(repo.index.diff("HEAD"))
|
||||
if staged_files:
|
||||
for item in staged_files[:25]: # Limit to first 25
|
||||
change_type = item.change_type
|
||||
status_lines.append(f"{change_type[0].upper()} {item.a_path}")
|
||||
if len(staged_files) > 25:
|
||||
status_lines.append(
|
||||
f"... and {len(staged_files) - 25} more staged files"
|
||||
)
|
||||
|
||||
# Check for unstaged changes
|
||||
unstaged_files = list(repo.index.diff(None))
|
||||
if unstaged_files:
|
||||
for item in unstaged_files[:25]: # Limit to first 25
|
||||
status_lines.append(f"M {item.a_path}")
|
||||
if len(unstaged_files) > 25:
|
||||
status_lines.append(
|
||||
f"... and {len(unstaged_files) - 25} more modified files"
|
||||
)
|
||||
|
||||
# Check for untracked files
|
||||
untracked_files = repo.untracked_files
|
||||
if untracked_files:
|
||||
for file in untracked_files[:25]: # Limit to first 25
|
||||
status_lines.append(f"?? {file}")
|
||||
if len(untracked_files) > 25:
|
||||
status_lines.append(
|
||||
f"... and {len(untracked_files) - 25} more untracked files"
|
||||
)
|
||||
|
||||
git_status = (
|
||||
"\n".join(status_lines) if status_lines else "Working tree clean"
|
||||
)
|
||||
|
||||
except Exception:
|
||||
git_status = "Unable to get git status"
|
||||
|
||||
# Get recent commits
|
||||
try:
|
||||
commits = []
|
||||
for commit in repo.iter_commits(max_count=5):
|
||||
short_hash = commit.hexsha[:7]
|
||||
message = commit.message.split("\n")[0] # First line only
|
||||
commits.append(f"{short_hash} {message}")
|
||||
recent_commits = "\n".join(commits)
|
||||
except Exception:
|
||||
recent_commits = "Unable to get recent commits"
|
||||
|
||||
return {
|
||||
"current_branch": current_branch,
|
||||
"main_branch": main_branch,
|
||||
"git_status": git_status,
|
||||
"recent_commits": recent_commits,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
"current_branch": None,
|
||||
"main_branch": None,
|
||||
"git_status": f"Error: {str(e)}",
|
||||
"recent_commits": f"Error: {str(e)}",
|
||||
}
|
||||
@@ -1,273 +0,0 @@
|
||||
"""MCP server implementing Hanzo capabilities."""
|
||||
|
||||
import os
|
||||
import atexit
|
||||
import signal
|
||||
import logging
|
||||
import warnings
|
||||
import threading
|
||||
from typing import Literal, cast, final
|
||||
|
||||
# Suppress litellm deprecation warnings about event loop
|
||||
warnings.filterwarnings(
|
||||
"ignore", message="There is no current event loop", category=DeprecationWarning
|
||||
)
|
||||
|
||||
try:
|
||||
from fastmcp import FastMCP
|
||||
except ImportError:
|
||||
# Fallback for older MCP versions
|
||||
from mcp.server import FastMCP
|
||||
|
||||
# Import our enhanced server
|
||||
from hanzo_mcp.tools import register_all_tools
|
||||
from hanzo_mcp.prompts import register_all_prompts
|
||||
from hanzo_mcp.server_enhanced import EnhancedFastMCP
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.shell.session_storage import SessionStorage
|
||||
|
||||
|
||||
@final
|
||||
class HanzoMCPServer:
|
||||
"""MCP server implementing Hanzo capabilities."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str = "hanzo",
|
||||
allowed_paths: list[str] | None = None,
|
||||
project_paths: list[str] | None = None,
|
||||
project_dir: str | None = None,
|
||||
mcp_instance: FastMCP | None = None,
|
||||
agent_model: str | None = None,
|
||||
agent_max_tokens: int | None = None,
|
||||
agent_api_key: str | None = None,
|
||||
agent_base_url: str | None = None,
|
||||
agent_max_iterations: int = 10,
|
||||
agent_max_tool_uses: int = 30,
|
||||
enable_agent_tool: bool = False,
|
||||
command_timeout: float = 120.0,
|
||||
disable_write_tools: bool = False,
|
||||
disable_search_tools: bool = False,
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 8888,
|
||||
enabled_tools: dict[str, bool] | None = None,
|
||||
disabled_tools: list[str] | None = None,
|
||||
):
|
||||
"""Initialize the Hanzo AI server.
|
||||
|
||||
Args:
|
||||
name: The name of the server
|
||||
allowed_paths: list of paths that the server is allowed to access
|
||||
project_paths: list of project paths to generate prompts for
|
||||
project_dir: single project directory (added to allowed_paths and project_paths)
|
||||
mcp_instance: Optional FastMCP instance for testing
|
||||
agent_model: Optional model name for agent tool in LiteLLM format
|
||||
agent_max_tokens: Optional maximum tokens for agent responses
|
||||
agent_api_key: Optional API key for the LLM provider
|
||||
agent_base_url: Optional base URL for the LLM provider API endpoint
|
||||
agent_max_iterations: Maximum number of iterations for agent (default: 10)
|
||||
agent_max_tool_uses: Maximum number of total tool uses for agent (default: 30)
|
||||
enable_agent_tool: Whether to enable the agent tool (default: False)
|
||||
command_timeout: Default timeout for command execution in seconds (default: 120.0)
|
||||
disable_write_tools: Whether to disable write tools (default: False)
|
||||
disable_search_tools: Whether to disable search tools (default: False)
|
||||
host: Host for SSE server (default: 127.0.0.1)
|
||||
port: Port for SSE server (default: 3000)
|
||||
enabled_tools: Dictionary of individual tool enable states (default: None)
|
||||
disabled_tools: List of tool names to disable (default: None)
|
||||
"""
|
||||
# Use enhanced server for automatic context normalization
|
||||
self.mcp = mcp_instance if mcp_instance is not None else EnhancedFastMCP(name)
|
||||
|
||||
# Initialize permissions and command executor
|
||||
self.permission_manager = PermissionManager()
|
||||
|
||||
# Handle project_dir parameter
|
||||
if project_dir:
|
||||
if allowed_paths is None:
|
||||
allowed_paths = []
|
||||
if project_dir not in allowed_paths:
|
||||
allowed_paths.append(project_dir)
|
||||
if project_paths is None:
|
||||
project_paths = []
|
||||
if project_dir not in project_paths:
|
||||
project_paths.append(project_dir)
|
||||
|
||||
# Add allowed paths
|
||||
if allowed_paths:
|
||||
for path in allowed_paths:
|
||||
self.permission_manager.add_allowed_path(path)
|
||||
|
||||
# Store paths and options
|
||||
self.project_paths = project_paths
|
||||
self.project_dir = project_dir
|
||||
self.disable_write_tools = disable_write_tools
|
||||
self.disable_search_tools = disable_search_tools
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.enabled_tools = enabled_tools or {}
|
||||
self.disabled_tools = disabled_tools or []
|
||||
|
||||
# Store agent options
|
||||
self.agent_model = agent_model
|
||||
self.agent_max_tokens = agent_max_tokens
|
||||
self.agent_api_key = agent_api_key
|
||||
self.agent_base_url = agent_base_url
|
||||
self.agent_max_iterations = agent_max_iterations
|
||||
self.agent_max_tool_uses = agent_max_tool_uses
|
||||
self.enable_agent_tool = enable_agent_tool
|
||||
self.command_timeout = command_timeout
|
||||
|
||||
# Initialize cleanup tracking
|
||||
self._cleanup_thread: threading.Thread | None = None
|
||||
self._shutdown_event = threading.Event()
|
||||
self._cleanup_registered = False
|
||||
|
||||
# Apply disabled_tools to enabled_tools
|
||||
final_enabled_tools = self.enabled_tools.copy()
|
||||
for tool_name in self.disabled_tools:
|
||||
final_enabled_tools[tool_name] = False
|
||||
|
||||
# Store the final processed tool configuration
|
||||
self.enabled_tools = final_enabled_tools
|
||||
|
||||
# Register all tools
|
||||
register_all_tools(
|
||||
mcp_server=self.mcp,
|
||||
permission_manager=self.permission_manager,
|
||||
agent_model=self.agent_model,
|
||||
agent_max_tokens=self.agent_max_tokens,
|
||||
agent_api_key=self.agent_api_key,
|
||||
agent_base_url=self.agent_base_url,
|
||||
agent_max_iterations=self.agent_max_iterations,
|
||||
agent_max_tool_uses=self.agent_max_tool_uses,
|
||||
enable_agent_tool=self.enable_agent_tool,
|
||||
disable_write_tools=self.disable_write_tools,
|
||||
disable_search_tools=self.disable_search_tools,
|
||||
enabled_tools=final_enabled_tools,
|
||||
)
|
||||
|
||||
register_all_prompts(mcp_server=self.mcp, projects=self.project_paths)
|
||||
|
||||
def _setup_cleanup_handlers(self) -> None:
|
||||
"""Set up signal handlers and background cleanup thread."""
|
||||
if self._cleanup_registered:
|
||||
return
|
||||
|
||||
# Register cleanup on normal exit
|
||||
atexit.register(self._cleanup_sessions)
|
||||
|
||||
# Register signal handlers for graceful shutdown
|
||||
def signal_handler(signum, frame):
|
||||
import sys
|
||||
|
||||
# Only log if not stdio transport
|
||||
if hasattr(self, "_transport") and self._transport != "stdio":
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info("\nShutting down gracefully...")
|
||||
self._cleanup_sessions()
|
||||
self._shutdown_event.set()
|
||||
sys.exit(0)
|
||||
|
||||
signal.signal(signal.SIGTERM, signal_handler)
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
# Start background cleanup thread for periodic cleanup
|
||||
self._cleanup_thread = threading.Thread(
|
||||
target=self._background_cleanup, daemon=True
|
||||
)
|
||||
self._cleanup_thread.start()
|
||||
|
||||
self._cleanup_registered = True
|
||||
|
||||
def _background_cleanup(self) -> None:
|
||||
"""Background thread for periodic session cleanup."""
|
||||
while not self._shutdown_event.is_set():
|
||||
try:
|
||||
# Clean up expired sessions every 2 minutes
|
||||
# Using shorter TTL of 5 minutes (300 seconds)
|
||||
SessionStorage.cleanup_expired_sessions(max_age_seconds=300)
|
||||
|
||||
# Wait for 2 minutes or until shutdown
|
||||
self._shutdown_event.wait(timeout=120)
|
||||
except Exception:
|
||||
# Ignore cleanup errors and continue
|
||||
pass
|
||||
|
||||
def _cleanup_sessions(self) -> None:
|
||||
"""Clean up all active sessions."""
|
||||
try:
|
||||
cleared_count = SessionStorage.clear_all_sessions()
|
||||
if cleared_count > 0:
|
||||
# Only log if not stdio transport
|
||||
if hasattr(self, "_transport") and self._transport != "stdio":
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(f"Cleaned up {cleared_count} tmux sessions on shutdown")
|
||||
except Exception:
|
||||
# Ignore cleanup errors during shutdown
|
||||
pass
|
||||
|
||||
def run(self, transport: str = "stdio", allowed_paths: list[str] | None = None):
|
||||
"""Run the MCP server.
|
||||
|
||||
Args:
|
||||
transport: The transport to use (stdio or sse)
|
||||
allowed_paths: list of paths that the server is allowed to access
|
||||
"""
|
||||
# Store transport for later use
|
||||
self._transport = transport
|
||||
|
||||
# Add allowed paths if provided
|
||||
allowed_paths_list = allowed_paths or []
|
||||
for path in allowed_paths_list:
|
||||
self.permission_manager.add_allowed_path(path)
|
||||
|
||||
# Show compute nodes only in non-stdio mode (to avoid corrupting protocol)
|
||||
if transport != "stdio" and not os.environ.get("HANZO_QUIET"):
|
||||
try:
|
||||
from hanzo_mcp.compute_nodes import ComputeNodeDetector
|
||||
|
||||
detector = ComputeNodeDetector()
|
||||
summary = detector.get_node_summary()
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(f"🖥️ {summary}")
|
||||
except Exception:
|
||||
# Silently ignore if compute node detection fails
|
||||
pass
|
||||
|
||||
# Set up cleanup handlers before running
|
||||
self._setup_cleanup_handlers()
|
||||
|
||||
# Run the server
|
||||
transport_type = cast(Literal["stdio", "sse"], transport)
|
||||
self.mcp.run(transport=transport_type)
|
||||
|
||||
|
||||
def create_server(
|
||||
name: str = "hanzo-mcp",
|
||||
allowed_paths: list[str] | None = None,
|
||||
enable_all_tools: bool = False,
|
||||
**kwargs,
|
||||
) -> HanzoMCPServer:
|
||||
"""Create a Hanzo MCP server instance.
|
||||
|
||||
Args:
|
||||
name: Server name
|
||||
allowed_paths: List of allowed file paths
|
||||
enable_all_tools: Enable all tools including agent tools
|
||||
**kwargs: Additional server configuration
|
||||
|
||||
Returns:
|
||||
HanzoMCPServer instance
|
||||
"""
|
||||
if enable_all_tools:
|
||||
kwargs["enable_agent_tool"] = True
|
||||
|
||||
return HanzoMCPServer(name=name, allowed_paths=allowed_paths, **kwargs)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point for the server."""
|
||||
from hanzo_mcp.cli import main as cli_main
|
||||
|
||||
cli_main()
|
||||
@@ -1,62 +0,0 @@
|
||||
"""Enhanced MCP server with automatic context normalization.
|
||||
|
||||
This module provides an enhanced FastMCP server that automatically
|
||||
applies context normalization to all registered tools.
|
||||
"""
|
||||
|
||||
from typing import Callable
|
||||
|
||||
from mcp.server import FastMCP
|
||||
|
||||
from hanzo_mcp.tools.common.decorators import with_context_normalization
|
||||
|
||||
|
||||
class EnhancedFastMCP(FastMCP):
|
||||
"""Enhanced FastMCP server with automatic context normalization.
|
||||
|
||||
This server automatically wraps all tool registrations with context
|
||||
normalization, ensuring that tools work properly when called externally
|
||||
with serialized context parameters.
|
||||
"""
|
||||
|
||||
def tool(self, name: str | None = None, description: str | None = None) -> Callable:
|
||||
"""Enhanced tool decorator that includes automatic context normalization.
|
||||
|
||||
Args:
|
||||
name: Tool name (defaults to function name)
|
||||
description: Tool description
|
||||
|
||||
Returns:
|
||||
Decorator function that registers the tool with context normalization
|
||||
"""
|
||||
# Get the original decorator from parent class
|
||||
original_decorator = super().tool(name=name, description=description)
|
||||
|
||||
# Create our enhanced decorator
|
||||
def enhanced_decorator(func: Callable) -> Callable:
|
||||
# Apply context normalization first
|
||||
# Check if function has ctx parameter
|
||||
import inspect
|
||||
|
||||
sig = inspect.signature(func)
|
||||
if "ctx" in sig.parameters:
|
||||
normalized_func = with_context_normalization(func)
|
||||
else:
|
||||
normalized_func = func
|
||||
|
||||
# Then apply the original decorator
|
||||
return original_decorator(normalized_func)
|
||||
|
||||
return enhanced_decorator
|
||||
|
||||
|
||||
def create_enhanced_server(name: str = "hanzo") -> EnhancedFastMCP:
|
||||
"""Create an enhanced MCP server with automatic context normalization.
|
||||
|
||||
Args:
|
||||
name: Server name
|
||||
|
||||
Returns:
|
||||
Enhanced FastMCP server instance
|
||||
"""
|
||||
return EnhancedFastMCP(name)
|
||||
@@ -1,494 +0,0 @@
|
||||
"""Tools package for Hanzo AI.
|
||||
|
||||
This package contains all tools for the Hanzo MCP server. To keep imports
|
||||
robust across environments (e.g., older Python during CI), heavy imports are
|
||||
guarded. Submodules can still be imported directly, e.g.:
|
||||
|
||||
from hanzo_mcp.tools.llm.llm_tool import LLMTool
|
||||
"""
|
||||
|
||||
# Defer annotation evaluation to avoid import-time NameErrors in constrained envs
|
||||
from __future__ import annotations
|
||||
|
||||
# NOTE: Keep top-level imports resilient for environments without optional deps
|
||||
try: # pragma: no cover - import guards for CI environments
|
||||
from mcp.server import FastMCP # type: ignore
|
||||
except Exception: # pragma: no cover
|
||||
|
||||
class FastMCP: # type: ignore
|
||||
pass
|
||||
|
||||
|
||||
try: # pragma: no cover
|
||||
from hanzo_mcp.tools.llm import (
|
||||
LLMTool,
|
||||
ConsensusTool,
|
||||
LLMManageTool,
|
||||
create_provider_tools,
|
||||
)
|
||||
from hanzo_mcp.tools.mcp import (
|
||||
MCPTool,
|
||||
McpAddTool,
|
||||
McpStatsTool,
|
||||
McpRemoveTool,
|
||||
)
|
||||
from hanzo_mcp.tools.todo import register_todo_tools
|
||||
from hanzo_mcp.tools.agent import register_agent_tools
|
||||
from hanzo_mcp.tools.shell import register_shell_tools
|
||||
from hanzo_mcp.tools.editor import (
|
||||
NeovimEditTool,
|
||||
NeovimCommandTool,
|
||||
NeovimSessionTool,
|
||||
)
|
||||
from hanzo_mcp.tools.vector import register_vector_tools
|
||||
from hanzo_mcp.tools.jupyter import register_jupyter_tools
|
||||
from hanzo_mcp.tools.database import DatabaseManager, register_database_tools
|
||||
from hanzo_mcp.tools.filesystem import register_filesystem_tools
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.mode import activate_mode_from_env
|
||||
from hanzo_mcp.tools.common.stats import StatsTool
|
||||
from hanzo_mcp.tools.common.tool_list import ToolListTool
|
||||
from hanzo_mcp.tools.config.mode_tool import mode_tool
|
||||
from hanzo_mcp.tools.common.mode_loader import ModeLoader
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.common.tool_enable import ToolEnableTool
|
||||
from hanzo_mcp.tools.common.tool_disable import ToolDisableTool
|
||||
|
||||
# Try memory tools
|
||||
try:
|
||||
from hanzo_mcp.tools.memory import register_memory_tools
|
||||
|
||||
MEMORY_TOOLS_AVAILABLE = True
|
||||
except Exception:
|
||||
MEMORY_TOOLS_AVAILABLE = False
|
||||
register_memory_tools = None # type: ignore
|
||||
except Exception:
|
||||
# Minimal surface to allow submodule imports elsewhere
|
||||
pass
|
||||
|
||||
# Try to import LSP tool
|
||||
try:
|
||||
from hanzo_mcp.tools.lsp import LSPTool, create_lsp_tool
|
||||
|
||||
LSP_TOOL_AVAILABLE = True
|
||||
except ImportError:
|
||||
LSP_TOOL_AVAILABLE = False
|
||||
|
||||
|
||||
def register_all_tools(
|
||||
mcp_server: FastMCP,
|
||||
permission_manager: PermissionManager,
|
||||
agent_model: str | None = None,
|
||||
agent_max_tokens: int | None = None,
|
||||
agent_api_key: str | None = None,
|
||||
agent_base_url: str | None = None,
|
||||
agent_max_iterations: int = 10,
|
||||
agent_max_tool_uses: int = 30,
|
||||
enable_agent_tool: bool = False,
|
||||
disable_write_tools: bool = False,
|
||||
disable_search_tools: bool = False,
|
||||
enabled_tools: dict[str, bool] | None = None,
|
||||
vector_config: dict | None = None,
|
||||
use_mode: bool = True,
|
||||
force_mode: str | None = None,
|
||||
) -> list[BaseTool]:
|
||||
"""Register all Hanzo tools with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
permission_manager: Permission manager for access control
|
||||
agent_model: Optional model name for agent tool in LiteLLM format
|
||||
agent_max_tokens: Optional maximum tokens for agent responses
|
||||
agent_api_key: Optional API key for the LLM provider
|
||||
agent_base_url: Optional base URL for the LLM provider API endpoint
|
||||
agent_max_iterations: Maximum number of iterations for agent (default: 10)
|
||||
agent_max_tool_uses: Maximum number of total tool uses for agent (default: 30)
|
||||
enable_agent_tool: Whether to enable the agent tool (default: False)
|
||||
disable_write_tools: Whether to disable write tools (default: False)
|
||||
disable_search_tools: Whether to disable search tools (default: False)
|
||||
enabled_tools: Dictionary of individual tool enable/disable states (default: None)
|
||||
vector_config: Vector store configuration (default: None)
|
||||
use_mode: Whether to use mode system for tool configuration (default: True)
|
||||
force_mode: Force a specific mode to be active (default: None)
|
||||
|
||||
Returns:
|
||||
List of registered BaseTool instances
|
||||
"""
|
||||
# Dictionary to store all registered tools
|
||||
all_tools: dict[str, BaseTool] = {}
|
||||
|
||||
# Load user plugins early
|
||||
try:
|
||||
plugins = load_user_plugins()
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
if plugins:
|
||||
logger.info(
|
||||
f"Loaded {len(plugins)} user plugin tools: {', '.join(plugins.keys())}"
|
||||
)
|
||||
except Exception as e:
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.warning(f"Failed to load user plugins: {e}")
|
||||
plugins = {}
|
||||
|
||||
# Apply mode configuration if enabled
|
||||
if use_mode:
|
||||
# First check for mode activation from environment
|
||||
activate_mode_from_env()
|
||||
|
||||
tool_config = ModeLoader.get_enabled_tools_from_mode(
|
||||
base_enabled_tools=enabled_tools, force_mode=force_mode
|
||||
)
|
||||
# Apply mode environment variables
|
||||
ModeLoader.apply_environment_from_mode()
|
||||
else:
|
||||
# Use individual tool configuration if provided, otherwise fall back to category-level flags
|
||||
tool_config = enabled_tools or {}
|
||||
|
||||
def is_tool_enabled(tool_name: str, category_enabled: bool = True) -> bool:
|
||||
"""Check if a specific tool should be enabled."""
|
||||
if tool_name in tool_config:
|
||||
return tool_config[tool_name]
|
||||
return category_enabled
|
||||
|
||||
# Register filesystem tools with individual configuration
|
||||
filesystem_enabled = {
|
||||
"read": is_tool_enabled("read", True),
|
||||
"write": is_tool_enabled("write", not disable_write_tools),
|
||||
"edit": is_tool_enabled("edit", not disable_write_tools),
|
||||
"multi_edit": is_tool_enabled("multi_edit", not disable_write_tools),
|
||||
"directory_tree": is_tool_enabled("directory_tree", True),
|
||||
"ast": is_tool_enabled("ast", not disable_search_tools),
|
||||
"rules": is_tool_enabled("rules", True),
|
||||
"search": is_tool_enabled("search", not disable_search_tools),
|
||||
"find": is_tool_enabled("find", True), # Fast file finder
|
||||
}
|
||||
|
||||
# Vector tools setup (needed for search)
|
||||
project_manager = None
|
||||
vector_enabled = {
|
||||
"vector_index": is_tool_enabled("vector_index", False),
|
||||
"vector_search": is_tool_enabled("vector_search", False),
|
||||
}
|
||||
|
||||
# Create project manager if vector tools, batch_search, or unified_search are enabled
|
||||
if (
|
||||
any(vector_enabled.values())
|
||||
or filesystem_enabled.get("batch_search", False)
|
||||
or filesystem_enabled.get("search", False)
|
||||
):
|
||||
if vector_config:
|
||||
from hanzo_mcp.tools.vector.project_manager import ProjectVectorManager
|
||||
|
||||
search_paths = [str(path) for path in permission_manager.allowed_paths]
|
||||
project_manager = ProjectVectorManager(
|
||||
global_db_path=vector_config.get("data_path"),
|
||||
embedding_model=vector_config.get(
|
||||
"embedding_model", "text-embedding-3-small"
|
||||
),
|
||||
dimension=vector_config.get("dimension", 1536),
|
||||
)
|
||||
# Auto-detect projects from search paths
|
||||
detected_projects = project_manager.detect_projects(search_paths)
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(f"Detected {len(detected_projects)} projects with LLM.md files")
|
||||
|
||||
filesystem_tools = register_filesystem_tools(
|
||||
mcp_server,
|
||||
permission_manager,
|
||||
enabled_tools=filesystem_enabled,
|
||||
project_manager=project_manager,
|
||||
)
|
||||
for tool in filesystem_tools:
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register jupyter tools if enabled
|
||||
jupyter_enabled = {
|
||||
"jupyter": is_tool_enabled("jupyter", True), # Unified tool
|
||||
"notebook_read": is_tool_enabled("notebook_read", True), # Legacy support
|
||||
"notebook_edit": is_tool_enabled("notebook_edit", True), # Legacy support
|
||||
}
|
||||
|
||||
if any(jupyter_enabled.values()):
|
||||
jupyter_tools = register_jupyter_tools(
|
||||
mcp_server, permission_manager, enabled_tools=jupyter_enabled
|
||||
)
|
||||
for tool in jupyter_tools:
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register shell tools if enabled
|
||||
if is_tool_enabled("run_command", True):
|
||||
shell_tools = register_shell_tools(mcp_server, permission_manager)
|
||||
for tool in shell_tools:
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register agent tools if enabled
|
||||
agent_enabled = (
|
||||
enable_agent_tool
|
||||
or is_tool_enabled("agent", False)
|
||||
or is_tool_enabled("dispatch_agent", False)
|
||||
)
|
||||
swarm_enabled = is_tool_enabled("swarm", False)
|
||||
|
||||
if agent_enabled or swarm_enabled:
|
||||
agent_tools = register_agent_tools(
|
||||
mcp_server,
|
||||
permission_manager,
|
||||
agent_model=agent_model,
|
||||
agent_max_tokens=agent_max_tokens,
|
||||
agent_api_key=agent_api_key,
|
||||
agent_base_url=agent_base_url,
|
||||
agent_max_iterations=agent_max_iterations,
|
||||
agent_max_tool_uses=agent_max_tool_uses,
|
||||
)
|
||||
# Filter based on what's enabled
|
||||
for tool in agent_tools:
|
||||
if (
|
||||
tool.name == "agent"
|
||||
and agent_enabled
|
||||
or tool.name == "swarm"
|
||||
and swarm_enabled
|
||||
):
|
||||
all_tools[tool.name] = tool
|
||||
elif tool.name in ["claude", "codex", "gemini", "grok", "code_auth"]:
|
||||
# CLI tools and auth are always included when agent tools are enabled
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register todo tools if enabled
|
||||
todo_enabled = {
|
||||
"todo": is_tool_enabled("todo", True),
|
||||
# Backward compatibility - if old names are used, enable the unified tool
|
||||
"todo_read": is_tool_enabled("todo_read", True),
|
||||
"todo_write": is_tool_enabled("todo_write", True),
|
||||
}
|
||||
|
||||
# Enable unified todo if any of the todo tools are enabled
|
||||
if any(todo_enabled.values()):
|
||||
todo_tools = register_todo_tools(mcp_server, enabled_tools={"todo": True})
|
||||
for tool in todo_tools:
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register thinking tool if enabled
|
||||
if is_tool_enabled("think", True):
|
||||
thinking_tool = register_thinking_tool(mcp_server)
|
||||
for tool in thinking_tool:
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register critic tool if enabled
|
||||
if is_tool_enabled("critic", True):
|
||||
critic_tools = register_critic_tool(mcp_server)
|
||||
for tool in critic_tools:
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register vector tools if enabled (reuse project_manager if available)
|
||||
if any(vector_enabled.values()) and project_manager:
|
||||
vector_tools = register_vector_tools(
|
||||
mcp_server,
|
||||
permission_manager,
|
||||
vector_config=vector_config,
|
||||
enabled_tools=vector_enabled,
|
||||
project_manager=project_manager,
|
||||
)
|
||||
for tool in vector_tools:
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register batch tool if enabled (batch tool is typically always enabled)
|
||||
if is_tool_enabled("batch", True):
|
||||
register_batch_tool(mcp_server, all_tools)
|
||||
|
||||
# Register database tools if enabled
|
||||
db_manager = None
|
||||
database_enabled = {
|
||||
"sql_query": is_tool_enabled("sql_query", True),
|
||||
"sql_search": is_tool_enabled("sql_search", True),
|
||||
"sql_stats": is_tool_enabled("sql_stats", True),
|
||||
"graph_add": is_tool_enabled("graph_add", True),
|
||||
"graph_remove": is_tool_enabled("graph_remove", True),
|
||||
"graph_query": is_tool_enabled("graph_query", True),
|
||||
"graph_search": is_tool_enabled("graph_search", True),
|
||||
"graph_stats": is_tool_enabled("graph_stats", True),
|
||||
}
|
||||
|
||||
if any(database_enabled.values()):
|
||||
db_manager = DatabaseManager(permission_manager)
|
||||
database_tools = register_database_tools(
|
||||
mcp_server,
|
||||
permission_manager,
|
||||
db_manager=db_manager,
|
||||
)
|
||||
# Filter based on enabled state
|
||||
for tool in database_tools:
|
||||
if database_enabled.get(tool.name, True):
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register unified MCP tool if enabled
|
||||
if is_tool_enabled("mcp", True):
|
||||
tool = MCPTool()
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register legacy MCP tools if explicitly enabled (disabled by default)
|
||||
legacy_mcp_enabled = {
|
||||
"mcp_add": is_tool_enabled("mcp_add", False),
|
||||
"mcp_remove": is_tool_enabled("mcp_remove", False),
|
||||
"mcp_stats": is_tool_enabled("mcp_stats", False),
|
||||
}
|
||||
|
||||
if legacy_mcp_enabled.get("mcp_add", False):
|
||||
tool = McpAddTool()
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
if legacy_mcp_enabled.get("mcp_remove", False):
|
||||
tool = McpRemoveTool()
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
if legacy_mcp_enabled.get("mcp_stats", False):
|
||||
tool = McpStatsTool()
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register system tools (always enabled)
|
||||
# Tool enable/disable tools
|
||||
tool_enable = ToolEnableTool()
|
||||
tool_enable.register(mcp_server)
|
||||
all_tools[tool_enable.name] = tool_enable
|
||||
|
||||
tool_disable = ToolDisableTool()
|
||||
tool_disable.register(mcp_server)
|
||||
all_tools[tool_disable.name] = tool_disable
|
||||
|
||||
tool_list = ToolListTool()
|
||||
tool_list.register(mcp_server)
|
||||
all_tools[tool_list.name] = tool_list
|
||||
|
||||
# Stats tool
|
||||
stats_tool = StatsTool(db_manager=db_manager)
|
||||
stats_tool.register(mcp_server)
|
||||
all_tools[stats_tool.name] = stats_tool
|
||||
|
||||
# Mode tool (always enabled for managing tool sets)
|
||||
mode_tool.register(mcp_server)
|
||||
all_tools[mode_tool.name] = mode_tool
|
||||
|
||||
# Register editor tools if enabled
|
||||
editor_enabled = {
|
||||
"neovim_edit": is_tool_enabled("neovim_edit", True),
|
||||
"neovim_command": is_tool_enabled("neovim_command", True),
|
||||
"neovim_session": is_tool_enabled("neovim_session", True),
|
||||
}
|
||||
|
||||
if editor_enabled.get("neovim_edit", True):
|
||||
tool = NeovimEditTool(permission_manager)
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
if editor_enabled.get("neovim_command", True):
|
||||
tool = NeovimCommandTool(permission_manager)
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
if editor_enabled.get("neovim_session", True):
|
||||
tool = NeovimSessionTool()
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register unified LLM tool if enabled
|
||||
if is_tool_enabled("llm", True):
|
||||
tool = LLMTool()
|
||||
if tool.available_providers: # Only register if API keys found
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register consensus tool if enabled (enabled by default)
|
||||
if is_tool_enabled("consensus", True):
|
||||
tool = ConsensusTool()
|
||||
if tool.llm_tool.available_providers:
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register legacy LLM tools if explicitly enabled (disabled by default)
|
||||
legacy_llm_enabled = {
|
||||
"llm_legacy": is_tool_enabled("llm_legacy", False),
|
||||
"consensus": is_tool_enabled("consensus", False),
|
||||
"llm_manage": is_tool_enabled("llm_manage", False),
|
||||
}
|
||||
|
||||
if legacy_llm_enabled.get("llm_legacy", False):
|
||||
tool = LLMTool()
|
||||
if tool.available_providers:
|
||||
tool.register(mcp_server)
|
||||
all_tools["llm_legacy"] = tool
|
||||
|
||||
if legacy_llm_enabled.get("consensus", False):
|
||||
tool = ConsensusTool()
|
||||
if tool.llm_tool.available_providers:
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
if legacy_llm_enabled.get("llm_manage", False):
|
||||
tool = LLMManageTool()
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register provider-specific LLM tools (disabled by default)
|
||||
if is_tool_enabled("provider_specific_llm", False):
|
||||
provider_tools = create_provider_tools()
|
||||
for tool in provider_tools:
|
||||
if is_tool_enabled(tool.name, False):
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
# Register memory tools if enabled
|
||||
memory_enabled = {
|
||||
"recall_memories": is_tool_enabled("recall_memories", True),
|
||||
"create_memories": is_tool_enabled("create_memories", True),
|
||||
"update_memories": is_tool_enabled("update_memories", True),
|
||||
"delete_memories": is_tool_enabled("delete_memories", True),
|
||||
"manage_memories": is_tool_enabled("manage_memories", True),
|
||||
"recall_facts": is_tool_enabled("recall_facts", True),
|
||||
"store_facts": is_tool_enabled("store_facts", True),
|
||||
"summarize_to_memory": is_tool_enabled("summarize_to_memory", True),
|
||||
"manage_knowledge_bases": is_tool_enabled("manage_knowledge_bases", True),
|
||||
}
|
||||
|
||||
if any(memory_enabled.values()) and MEMORY_TOOLS_AVAILABLE:
|
||||
try:
|
||||
memory_tools = register_memory_tools(
|
||||
mcp_server, permission_manager, user_id="default", project_id="default"
|
||||
)
|
||||
# Filter based on enabled state
|
||||
for tool in memory_tools:
|
||||
if memory_enabled.get(tool.name, True):
|
||||
all_tools[tool.name] = tool
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to register memory tools: {e}")
|
||||
|
||||
# Register LSP tool if enabled
|
||||
if is_tool_enabled("lsp", True) and LSP_TOOL_AVAILABLE:
|
||||
try:
|
||||
tool = create_lsp_tool()
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to register LSP tool: {e}")
|
||||
|
||||
# Register user plugins last (so they can override built-in tools)
|
||||
for plugin_name, plugin in plugins.items():
|
||||
if is_tool_enabled(plugin_name, True):
|
||||
try:
|
||||
tool = plugin.tool_class()
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
logger.info(f"Registered plugin tool: {plugin_name}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to register plugin tool {plugin_name}: {e}")
|
||||
|
||||
# Return all registered tools
|
||||
return list(all_tools.values())
|
||||
@@ -1,132 +0,0 @@
|
||||
"""Agent tools for Hanzo AI.
|
||||
|
||||
This module provides tools that allow Claude to delegate tasks to sub-agents,
|
||||
enabling concurrent execution of multiple operations and specialized processing.
|
||||
"""
|
||||
|
||||
from mcp.server import FastMCP
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool, ToolRegistry
|
||||
|
||||
# Import unified CLI tools (single source of truth)
|
||||
from hanzo_mcp.tools.agent.cli_tools import (
|
||||
GrokCLITool,
|
||||
AiderCLITool,
|
||||
ClineCLITool,
|
||||
CodexCLITool,
|
||||
ClaudeCLITool,
|
||||
GeminiCLITool,
|
||||
HanzoDevCLITool,
|
||||
OpenHandsCLITool,
|
||||
ClaudeCodeCLITool, # cc alias
|
||||
OpenHandsShortCLITool, # oh alias
|
||||
register_cli_tools,
|
||||
)
|
||||
|
||||
# Import the main implementations (using hanzo-agents SDK)
|
||||
from hanzo_mcp.tools.agent.agent_tool import AgentTool
|
||||
from hanzo_mcp.tools.agent.network_tool import NetworkTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.agent.code_auth_tool import CodeAuthTool
|
||||
|
||||
|
||||
def register_agent_tools(
|
||||
mcp_server: FastMCP,
|
||||
permission_manager: PermissionManager,
|
||||
agent_model: str | None = None,
|
||||
agent_max_tokens: int | None = None,
|
||||
agent_api_key: str | None = None,
|
||||
agent_base_url: str | None = None,
|
||||
agent_max_iterations: int = 10,
|
||||
agent_max_tool_uses: int = 30,
|
||||
) -> list[BaseTool]:
|
||||
"""Register agent tools with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
|
||||
permission_manager: Permission manager for access control
|
||||
agent_model: Optional model name for agent tool in LiteLLM format
|
||||
agent_max_tokens: Optional maximum tokens for agent responses
|
||||
agent_api_key: Optional API key for the LLM provider
|
||||
agent_base_url: Optional base URL for the LLM provider API endpoint
|
||||
agent_max_iterations: Maximum number of iterations for agent (default: 10)
|
||||
agent_max_tool_uses: Maximum number of total tool uses for agent (default: 30)
|
||||
|
||||
Returns:
|
||||
List of registered tools
|
||||
"""
|
||||
# Create agent tool
|
||||
agent_tool = AgentTool(
|
||||
permission_manager=permission_manager,
|
||||
model=agent_model,
|
||||
api_key=agent_api_key,
|
||||
base_url=agent_base_url,
|
||||
max_tokens=agent_max_tokens,
|
||||
max_iterations=agent_max_iterations,
|
||||
max_tool_uses=agent_max_tool_uses,
|
||||
)
|
||||
|
||||
# Register a swarm alias that forwards to AgentTool with default concurrency
|
||||
class SwarmAliasTool(BaseTool):
|
||||
name = "swarm"
|
||||
description = (
|
||||
"Alias for agent with concurrency. swarm == agent:5 by default.\n"
|
||||
"Use 'swarm' for parallel multi-agent runs; 'swarm:N' for N agents."
|
||||
)
|
||||
|
||||
def __init__(self, agent_tool: AgentTool):
|
||||
self._agent = agent_tool
|
||||
|
||||
async def call(self, ctx, **params): # type: ignore[override]
|
||||
# Default to 5 agents unless explicitly provided
|
||||
params = dict(params)
|
||||
params.setdefault("concurrency", 5)
|
||||
return await self._agent.call(ctx, **params)
|
||||
|
||||
def register(self, mcp_server: FastMCP): # type: ignore[override]
|
||||
tool_self = self
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def swarm(
|
||||
ctx,
|
||||
prompts: str | list[str], # forwarded
|
||||
concurrency: int | None = None,
|
||||
model: str | None = None,
|
||||
use_memory: bool | None = None,
|
||||
memory_backend: str | None = None,
|
||||
) -> str:
|
||||
p = {
|
||||
"prompts": prompts,
|
||||
}
|
||||
if concurrency is not None:
|
||||
p["concurrency"] = concurrency
|
||||
if model is not None:
|
||||
p["model"] = model
|
||||
if use_memory is not None:
|
||||
p["use_memory"] = use_memory
|
||||
if memory_backend is not None:
|
||||
p["memory_backend"] = memory_backend
|
||||
return await tool_self.call(ctx, **p)
|
||||
return tool_self
|
||||
|
||||
# Create auth management tool
|
||||
code_auth_tool = CodeAuthTool()
|
||||
|
||||
# Create network tool
|
||||
network_tool = NetworkTool(
|
||||
permission_manager=permission_manager,
|
||||
default_mode="hybrid", # Prefer local, fallback to cloud
|
||||
)
|
||||
|
||||
# Register core agent tools
|
||||
ToolRegistry.register_tool(mcp_server, agent_tool)
|
||||
ToolRegistry.register_tool(mcp_server, SwarmAliasTool(agent_tool))
|
||||
ToolRegistry.register_tool(mcp_server, network_tool)
|
||||
ToolRegistry.register_tool(mcp_server, code_auth_tool)
|
||||
|
||||
# Register all CLI tools (includes claude, codex, gemini, grok, etc.)
|
||||
cli_tools = register_cli_tools(mcp_server, permission_manager)
|
||||
|
||||
# Return list of registered tools
|
||||
return [agent_tool, network_tool, code_auth_tool] + cli_tools
|
||||
@@ -1,422 +0,0 @@
|
||||
"""Unified agent tool implementation.
|
||||
|
||||
This module provides the AgentTool for delegating tasks to sub-agents,
|
||||
supporting both one-off and long-running RPC modes, including A2A communication.
|
||||
"""
|
||||
|
||||
import re
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
import asyncio
|
||||
|
||||
# Import litellm with warnings suppressed
|
||||
import warnings
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
List,
|
||||
Unpack,
|
||||
Optional,
|
||||
Annotated,
|
||||
TypedDict,
|
||||
final,
|
||||
override,
|
||||
)
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", DeprecationWarning)
|
||||
from pydantic import Field
|
||||
from openai.types.chat import ChatCompletionMessageParam
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.jupyter import get_read_only_jupyter_tools
|
||||
from hanzo_mcp.tools.filesystem import get_read_only_filesystem_tools
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.agent.prompt import (
|
||||
get_default_model,
|
||||
get_system_prompt,
|
||||
get_allowed_agent_tools,
|
||||
)
|
||||
from hanzo_mcp.tools.common.context import (
|
||||
ToolContext,
|
||||
create_tool_context,
|
||||
)
|
||||
from hanzo_mcp.tools.common.batch_tool import BatchTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
# Parameter types
|
||||
Action = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Action: run (default), start, call, stop, list",
|
||||
default="run",
|
||||
),
|
||||
]
|
||||
|
||||
Prompts = Annotated[
|
||||
Optional[str | List[str]],
|
||||
Field(
|
||||
description="Task(s) for agent (must include absolute paths starting with /)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Mode = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Execution mode: oneoff (default) or rpc",
|
||||
default="oneoff",
|
||||
),
|
||||
]
|
||||
|
||||
AgentId = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Agent ID for RPC mode",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Method = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Method to call on RPC agent",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Args = Annotated[
|
||||
Optional[Dict[str, Any]],
|
||||
Field(
|
||||
description="Arguments for RPC method call",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Model = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Model to use (e.g., lm-studio/local-model, openai/gpt-4o)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class AgentParams(TypedDict, total=False):
|
||||
"""Parameters for agent tool."""
|
||||
|
||||
action: str
|
||||
prompts: Optional[str | List[str]]
|
||||
mode: str
|
||||
agent_id: Optional[str]
|
||||
method: Optional[str]
|
||||
args: Optional[Dict[str, Any]]
|
||||
model: Optional[str]
|
||||
|
||||
|
||||
class RPCAgent:
|
||||
"""Long-running RPC agent."""
|
||||
|
||||
def __init__(
|
||||
self, agent_id: str, model: str, system_prompt: str, tools: List[BaseTool]
|
||||
):
|
||||
self.agent_id = agent_id
|
||||
self.model = model
|
||||
self.system_prompt = system_prompt
|
||||
self.tools = tools
|
||||
self.messages: List[ChatCompletionMessageParam] = [
|
||||
{"role": "system", "content": system_prompt}
|
||||
]
|
||||
self.created_at = time.time()
|
||||
self.last_used = time.time()
|
||||
self.call_count = 0
|
||||
|
||||
async def call_method(
|
||||
self, method: str, args: Dict[str, Any], tool_ctx: ToolContext
|
||||
) -> str:
|
||||
"""Call a method on the RPC agent."""
|
||||
self.last_used = time.time()
|
||||
self.call_count += 1
|
||||
|
||||
# Build prompt based on method
|
||||
if method == "search":
|
||||
prompt = f"Search for: {args.get('query', 'unknown')}"
|
||||
elif method == "analyze":
|
||||
prompt = f"Analyze: {args.get('target', 'unknown')}"
|
||||
elif method == "execute":
|
||||
prompt = f"Execute: {args.get('command', 'unknown')}"
|
||||
else:
|
||||
# Generic method call
|
||||
prompt = f"Method: {method}, Args: {json.dumps(args)}"
|
||||
|
||||
# Add to conversation
|
||||
self.messages.append({"role": "user", "content": prompt})
|
||||
|
||||
# Get response
|
||||
# (simplified - would integrate with full agent execution logic)
|
||||
response = f"Executed {method} with args {args}"
|
||||
self.messages.append({"role": "assistant", "content": response})
|
||||
|
||||
return response
|
||||
|
||||
|
||||
@final
|
||||
class AgentTool(BaseTool):
|
||||
"""Unified agent tool with one-off and RPC modes."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
model: str | None = None,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
max_iterations: int = 10,
|
||||
max_tool_uses: int = 30,
|
||||
):
|
||||
"""Initialize the agent tool."""
|
||||
self.permission_manager = permission_manager
|
||||
self.model_override = model
|
||||
self.api_key_override = api_key
|
||||
self.base_url_override = base_url
|
||||
self.max_tokens_override = max_tokens
|
||||
self.max_iterations = max_iterations
|
||||
self.max_tool_uses = max_tool_uses
|
||||
|
||||
# RPC agent registry
|
||||
self._rpc_agents: Dict[str, RPCAgent] = {}
|
||||
|
||||
# Available tools
|
||||
self.available_tools: list[BaseTool] = []
|
||||
self.available_tools.extend(
|
||||
get_read_only_filesystem_tools(self.permission_manager)
|
||||
)
|
||||
self.available_tools.extend(
|
||||
get_read_only_jupyter_tools(self.permission_manager)
|
||||
)
|
||||
self.available_tools.append(
|
||||
BatchTool({t.name: t for t in self.available_tools})
|
||||
)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "agent"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
tools = [t.name for t in self.available_tools]
|
||||
|
||||
return f"""AI agents with tools: {", ".join(tools)}. Actions: run (default), start, call, stop, list.
|
||||
|
||||
Usage:
|
||||
agent "Search for config files in /project"
|
||||
agent --action start --mode rpc --model lm-studio/local-model
|
||||
agent --action call --agent-id abc123 --method search --args '{{"query": "database"}}'
|
||||
agent --action list
|
||||
|
||||
Modes:
|
||||
- oneoff: Single task execution (default)
|
||||
- rpc: Long-running agent for multiple calls (A2A support)"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[AgentParams],
|
||||
) -> str:
|
||||
"""Execute agent operation."""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract action
|
||||
action = params.get("action", "run")
|
||||
|
||||
# Route to appropriate handler
|
||||
if action == "run":
|
||||
return await self._handle_run(params, tool_ctx)
|
||||
elif action == "start":
|
||||
return await self._handle_start(params, tool_ctx)
|
||||
elif action == "call":
|
||||
return await self._handle_call(params, tool_ctx)
|
||||
elif action == "stop":
|
||||
return await self._handle_stop(params, tool_ctx)
|
||||
elif action == "list":
|
||||
return await self._handle_list(tool_ctx)
|
||||
else:
|
||||
return f"Error: Unknown action '{action}'. Valid actions: run, start, call, stop, list"
|
||||
|
||||
async def _handle_run(self, params: Dict[str, Any], tool_ctx: ToolContext) -> str:
|
||||
"""Handle one-off agent run (default action)."""
|
||||
prompts = params.get("prompts")
|
||||
if not prompts:
|
||||
return "Error: prompts required for run action"
|
||||
|
||||
# Convert to list
|
||||
if isinstance(prompts, str):
|
||||
prompt_list = [prompts]
|
||||
else:
|
||||
prompt_list = prompts
|
||||
|
||||
# Validate prompts
|
||||
for prompt in prompt_list:
|
||||
if not self._validate_prompt(prompt):
|
||||
return f"Error: Prompt must contain absolute paths starting with /: {prompt[:50]}..."
|
||||
|
||||
# Execute agents
|
||||
start_time = time.time()
|
||||
|
||||
if len(prompt_list) == 1:
|
||||
await tool_ctx.info("Launching agent")
|
||||
result = await self._execute_agent(
|
||||
prompt_list[0], params.get("model"), tool_ctx
|
||||
)
|
||||
else:
|
||||
await tool_ctx.info(f"Launching {len(prompt_list)} agents in parallel")
|
||||
result = await self._execute_multiple_agents(
|
||||
prompt_list, params.get("model"), tool_ctx
|
||||
)
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
return f"""Agent execution completed in {execution_time:.2f} seconds.
|
||||
|
||||
AGENT RESPONSE:
|
||||
{result}"""
|
||||
|
||||
async def _handle_start(self, params: Dict[str, Any], tool_ctx: ToolContext) -> str:
|
||||
"""Start a new RPC agent."""
|
||||
mode = params.get("mode", "oneoff")
|
||||
if mode != "rpc":
|
||||
return "Error: start action only valid for rpc mode"
|
||||
|
||||
# Generate agent ID
|
||||
agent_id = str(uuid.uuid4())[:8]
|
||||
|
||||
# Get model
|
||||
model = params.get("model") or get_default_model(self.model_override)
|
||||
|
||||
# Get available tools
|
||||
agent_tools = get_allowed_agent_tools(
|
||||
self.available_tools,
|
||||
self.permission_manager,
|
||||
)
|
||||
|
||||
# Create system prompt
|
||||
system_prompt = get_system_prompt(
|
||||
agent_tools,
|
||||
self.permission_manager,
|
||||
)
|
||||
|
||||
# Create RPC agent
|
||||
agent = RPCAgent(agent_id, model, system_prompt, agent_tools)
|
||||
self._rpc_agents[agent_id] = agent
|
||||
|
||||
await tool_ctx.info(f"Started RPC agent {agent_id} with model {model}")
|
||||
|
||||
return f"""Started RPC agent:
|
||||
- ID: {agent_id}
|
||||
- Model: {model}
|
||||
- Tools: {len(agent_tools)}
|
||||
|
||||
Use 'agent --action call --agent-id {agent_id} --method <method> --args <args>' to interact."""
|
||||
|
||||
async def _handle_call(self, params: Dict[str, Any], tool_ctx: ToolContext) -> str:
|
||||
"""Call method on RPC agent."""
|
||||
agent_id = params.get("agent_id")
|
||||
if not agent_id:
|
||||
return "Error: agent_id required for call action"
|
||||
|
||||
if agent_id not in self._rpc_agents:
|
||||
return f"Error: Agent {agent_id} not found. Use 'agent --action list' to see active agents."
|
||||
|
||||
method = params.get("method")
|
||||
if not method:
|
||||
return "Error: method required for call action"
|
||||
|
||||
args = params.get("args", {})
|
||||
|
||||
# Call agent method
|
||||
agent = self._rpc_agents[agent_id]
|
||||
await tool_ctx.info(f"Calling {method} on agent {agent_id}")
|
||||
|
||||
try:
|
||||
result = await agent.call_method(method, args, tool_ctx)
|
||||
return f"Agent {agent_id} response:\n{result}"
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Error calling agent: {str(e)}")
|
||||
return f"Error calling agent: {str(e)}"
|
||||
|
||||
async def _handle_stop(self, params: Dict[str, Any], tool_ctx: ToolContext) -> str:
|
||||
"""Stop an RPC agent."""
|
||||
agent_id = params.get("agent_id")
|
||||
if not agent_id:
|
||||
return "Error: agent_id required for stop action"
|
||||
|
||||
if agent_id not in self._rpc_agents:
|
||||
return f"Error: Agent {agent_id} not found"
|
||||
|
||||
agent = self._rpc_agents.pop(agent_id)
|
||||
await tool_ctx.info(f"Stopped agent {agent_id}")
|
||||
|
||||
return f"""Stopped agent {agent_id}:
|
||||
- Runtime: {time.time() - agent.created_at:.2f} seconds
|
||||
- Calls: {agent.call_count}"""
|
||||
|
||||
async def _handle_list(self, tool_ctx: ToolContext) -> str:
|
||||
"""List active RPC agents."""
|
||||
if not self._rpc_agents:
|
||||
return "No active RPC agents"
|
||||
|
||||
output = ["=== Active RPC Agents ==="]
|
||||
for agent_id, agent in self._rpc_agents.items():
|
||||
runtime = time.time() - agent.created_at
|
||||
idle = time.time() - agent.last_used
|
||||
output.append(f"\nAgent {agent_id}:")
|
||||
output.append(f" Model: {agent.model}")
|
||||
output.append(f" Runtime: {runtime:.2f}s")
|
||||
output.append(f" Idle: {idle:.2f}s")
|
||||
output.append(f" Calls: {agent.call_count}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _validate_prompt(self, prompt: str) -> bool:
|
||||
"""Validate that prompt contains absolute paths."""
|
||||
absolute_path_pattern = r"/(?:[^/\s]+/)*[^/\s]+"
|
||||
return bool(re.search(absolute_path_pattern, prompt))
|
||||
|
||||
async def _execute_agent(
|
||||
self, prompt: str, model: Optional[str], tool_ctx: ToolContext
|
||||
) -> str:
|
||||
"""Execute a single agent (simplified - would use full logic from agent_tool.py)."""
|
||||
# This would integrate the full agent execution logic from agent_tool.py
|
||||
# For now, return a placeholder
|
||||
return f"Executed agent with prompt: {prompt[:100]}..."
|
||||
|
||||
async def _execute_multiple_agents(
|
||||
self, prompts: List[str], model: Optional[str], tool_ctx: ToolContext
|
||||
) -> str:
|
||||
"""Execute multiple agents in parallel."""
|
||||
tasks = []
|
||||
for prompt in prompts:
|
||||
task = self._execute_agent(prompt, model, tool_ctx)
|
||||
tasks.append(task)
|
||||
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
formatted_results = []
|
||||
for i, result in enumerate(results):
|
||||
if isinstance(result, Exception):
|
||||
formatted_results.append(f"Agent {i + 1} Error:\n{str(result)}")
|
||||
else:
|
||||
formatted_results.append(f"Agent {i + 1} Result:\n{result}")
|
||||
|
||||
return "\n\n---\n\n".join(formatted_results)
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,522 +0,0 @@
|
||||
"""Agent tool implementation using hanzo-agents SDK.
|
||||
|
||||
This module implements the AgentTool that leverages the hanzo-agents SDK
|
||||
for sophisticated agent orchestration and execution.
|
||||
"""
|
||||
|
||||
import re
|
||||
import time
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
List,
|
||||
Unpack,
|
||||
Optional,
|
||||
TypedDict,
|
||||
final,
|
||||
override,
|
||||
)
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
# Import hanzo-agents SDK
|
||||
try:
|
||||
from hanzo_agents import (
|
||||
Tool,
|
||||
Agent,
|
||||
State,
|
||||
History,
|
||||
Network,
|
||||
InferenceResult,
|
||||
create_memory_kv,
|
||||
sequential_router,
|
||||
create_memory_vector,
|
||||
)
|
||||
from hanzo_agents.core.cli_agent import (
|
||||
GrokAgent,
|
||||
GeminiAgent,
|
||||
ClaudeCodeAgent,
|
||||
OpenAICodexAgent,
|
||||
)
|
||||
|
||||
HANZO_AGENTS_AVAILABLE = True
|
||||
except ImportError:
|
||||
HANZO_AGENTS_AVAILABLE = False
|
||||
|
||||
# Define stub classes when hanzo-agents is not available
|
||||
class State:
|
||||
"""Stub State class when hanzo-agents is not available."""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def to_dict(self):
|
||||
return {}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data):
|
||||
return cls()
|
||||
|
||||
class Tool:
|
||||
"""Stub Tool class when hanzo-agents is not available."""
|
||||
|
||||
pass
|
||||
|
||||
class Agent:
|
||||
"""Stub Agent class when hanzo-agents is not available."""
|
||||
|
||||
pass
|
||||
|
||||
class Network:
|
||||
"""Stub Network class when hanzo-agents is not available."""
|
||||
|
||||
pass
|
||||
|
||||
class History:
|
||||
"""Stub History class when hanzo-agents is not available."""
|
||||
|
||||
pass
|
||||
|
||||
class InferenceResult:
|
||||
"""Stub InferenceResult class when hanzo-agents is not available."""
|
||||
|
||||
def __init__(self, agent=None, content=None, metadata=None):
|
||||
self.agent = agent
|
||||
self.content = content
|
||||
self.metadata = metadata or {}
|
||||
|
||||
|
||||
from hanzo_mcp.tools.jupyter import get_read_only_jupyter_tools
|
||||
from hanzo_mcp.tools.filesystem import Edit, MultiEdit, get_read_only_filesystem_tools
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.agent.critic_tool import CriticTool
|
||||
from hanzo_mcp.tools.agent.iching_tool import IChingTool
|
||||
from hanzo_mcp.tools.agent.review_tool import ReviewTool
|
||||
from hanzo_mcp.tools.common.batch_tool import BatchTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.agent.clarification_tool import ClarificationTool
|
||||
from hanzo_mcp.tools.agent.clarification_protocol import (
|
||||
AgentClarificationMixin,
|
||||
)
|
||||
|
||||
|
||||
class AgentToolParams(TypedDict, total=False):
|
||||
"""Parameters for the AgentTool."""
|
||||
|
||||
prompts: str | list[str]
|
||||
model: Optional[str]
|
||||
use_memory: Optional[bool]
|
||||
memory_backend: Optional[str]
|
||||
concurrency: Optional[int]
|
||||
|
||||
|
||||
class MCPAgentState(State):
|
||||
"""State for MCP agents."""
|
||||
|
||||
def __init__(self, prompts: List[str], context: Dict[str, Any]):
|
||||
"""Initialize agent state."""
|
||||
super().__init__()
|
||||
self.prompts = prompts
|
||||
self.context = context
|
||||
self.current_prompt_index = 0
|
||||
self.results = []
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
base_dict = super().to_dict()
|
||||
base_dict.update(
|
||||
{
|
||||
"prompts": self.prompts,
|
||||
"context": self.context,
|
||||
"current_prompt_index": self.current_prompt_index,
|
||||
"results": self.results,
|
||||
}
|
||||
)
|
||||
return base_dict
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "MCPAgentState":
|
||||
"""Create from dictionary."""
|
||||
state = cls(prompts=data.get("prompts", []), context=data.get("context", {}))
|
||||
state.current_prompt_index = data.get("current_prompt_index", 0)
|
||||
state.results = data.get("results", [])
|
||||
for k, v in data.items():
|
||||
if k not in ["prompts", "context", "current_prompt_index", "results"]:
|
||||
state[k] = v
|
||||
return state
|
||||
|
||||
|
||||
class MCPToolAdapter(Tool):
|
||||
"""Adapter to wrap MCP tools for hanzo-agents."""
|
||||
|
||||
def __init__(self, mcp_tool: BaseTool, ctx: MCPContext):
|
||||
"""Initialize adapter."""
|
||||
self.mcp_tool = mcp_tool
|
||||
self.ctx = ctx
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""Get tool name."""
|
||||
return self.mcp_tool.name
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
"""Get tool description."""
|
||||
return self.mcp_tool.description
|
||||
|
||||
async def execute(self, **kwargs) -> str:
|
||||
"""Execute the MCP tool."""
|
||||
return await self.mcp_tool.call(self.ctx, **kwargs)
|
||||
|
||||
|
||||
class MCPAgent(Agent):
|
||||
"""Agent that executes MCP tasks."""
|
||||
|
||||
name = "mcp_agent"
|
||||
description = "Agent for executing MCP tasks"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
available_tools: List[BaseTool],
|
||||
permission_manager: PermissionManager,
|
||||
ctx: MCPContext,
|
||||
model: str = "model://anthropic/claude-3-5-sonnet-20241022",
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize MCP agent."""
|
||||
super().__init__(model=model, **kwargs)
|
||||
|
||||
self.available_tools = available_tools
|
||||
self.permission_manager = permission_manager
|
||||
self.ctx = ctx
|
||||
|
||||
# Register MCP tools as agent tools
|
||||
for mcp_tool in available_tools:
|
||||
adapter = MCPToolAdapter(mcp_tool, ctx)
|
||||
self.register_tool(adapter)
|
||||
|
||||
async def run(
|
||||
self, state: MCPAgentState, history: History, network: Network
|
||||
) -> InferenceResult:
|
||||
"""Execute the agent."""
|
||||
# Get current prompt
|
||||
if state.current_prompt_index >= len(state.prompts):
|
||||
return InferenceResult(
|
||||
agent=self.name,
|
||||
content="All prompts completed",
|
||||
metadata={"completed": True},
|
||||
)
|
||||
|
||||
prompt = state.prompts[state.current_prompt_index]
|
||||
|
||||
# Execute with tools
|
||||
messages = [
|
||||
{"role": "system", "content": self._get_system_prompt()},
|
||||
{"role": "user", "content": prompt},
|
||||
]
|
||||
|
||||
# Add history context
|
||||
for entry in history[-10:]:
|
||||
if entry.role == "assistant":
|
||||
messages.append({"role": "assistant", "content": entry.content})
|
||||
elif entry.role == "user":
|
||||
messages.append({"role": "user", "content": entry.content})
|
||||
|
||||
# Call model
|
||||
from hanzo_agents import ModelRegistry
|
||||
|
||||
adapter = ModelRegistry.get_adapter(self.model)
|
||||
response = await adapter.chat(messages)
|
||||
|
||||
# Update state
|
||||
state.current_prompt_index += 1
|
||||
state.results.append(response)
|
||||
|
||||
# Return result
|
||||
return InferenceResult(
|
||||
agent=self.name,
|
||||
content=response,
|
||||
metadata={
|
||||
"prompt_index": state.current_prompt_index - 1,
|
||||
"total_prompts": len(state.prompts),
|
||||
},
|
||||
)
|
||||
|
||||
def _get_system_prompt(self) -> str:
|
||||
"""Get system prompt for the agent."""
|
||||
tool_descriptions = []
|
||||
for tool in self.tools.values():
|
||||
tool_descriptions.append(f"- {tool.name}: {tool.description}")
|
||||
|
||||
return f"""You are an AI assistant with access to the following tools:
|
||||
|
||||
{chr(10).join(tool_descriptions)}
|
||||
|
||||
When you need to use a tool, respond with:
|
||||
TOOL: tool_name(arg1="value1", arg2="value2")
|
||||
|
||||
Important guidelines:
|
||||
- Always include absolute paths starting with / when working with files
|
||||
- Be thorough in your searches and analysis
|
||||
- Provide clear, actionable results
|
||||
- Edit files when requested to make changes
|
||||
"""
|
||||
|
||||
|
||||
@final
|
||||
class AgentTool(AgentClarificationMixin, BaseTool):
|
||||
"""Tool for delegating tasks to sub-agents using hanzo-agents SDK."""
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "agent"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
if not HANZO_AGENTS_AVAILABLE:
|
||||
return "Agent tool (hanzo-agents SDK not available - using fallback)"
|
||||
|
||||
at = [t.name for t in self.available_tools]
|
||||
return f"""Launch a new agent that has access to the following tools: {at}.
|
||||
|
||||
When to use the Agent tool:
|
||||
- If you are searching for a keyword like "config" or "logger"
|
||||
- When you need to perform edits across multiple files
|
||||
- When you need to delegate complex file modification tasks
|
||||
|
||||
When NOT to use the Agent tool:
|
||||
- If you want to read a specific file path
|
||||
- If you are searching for a specific class definition
|
||||
- Writing code and running bash commands
|
||||
- Other tasks that are not related to searching
|
||||
|
||||
Usage notes:
|
||||
1. Launch multiple agents concurrently whenever possible
|
||||
2. Agent results are not visible to the user - summarize them
|
||||
3. Each agent invocation is stateless
|
||||
4. The agent's outputs should generally be trusted
|
||||
5. Clearly tell the agent whether you expect it to write code or just do research"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
model: str | None = None,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
max_iterations: int = 10,
|
||||
max_tool_uses: int = 30,
|
||||
) -> None:
|
||||
"""Initialize the agent tool."""
|
||||
self.permission_manager = permission_manager
|
||||
self.model_override = model
|
||||
self.api_key_override = api_key
|
||||
self.base_url_override = base_url
|
||||
self.max_tokens_override = max_tokens
|
||||
self.max_iterations = max_iterations
|
||||
self.max_tool_uses = max_tool_uses
|
||||
|
||||
# Set up available tools
|
||||
self.available_tools: list[BaseTool] = []
|
||||
self.available_tools.extend(
|
||||
get_read_only_filesystem_tools(self.permission_manager)
|
||||
)
|
||||
self.available_tools.extend(
|
||||
get_read_only_jupyter_tools(self.permission_manager)
|
||||
)
|
||||
|
||||
# Add edit tools
|
||||
self.available_tools.append(Edit(self.permission_manager))
|
||||
self.available_tools.append(MultiEdit(self.permission_manager))
|
||||
|
||||
# Add special tools
|
||||
self.available_tools.append(ClarificationTool())
|
||||
self.available_tools.append(CriticTool())
|
||||
self.available_tools.append(ReviewTool())
|
||||
self.available_tools.append(IChingTool())
|
||||
|
||||
self.available_tools.append(
|
||||
BatchTool({t.name: t for t in self.available_tools})
|
||||
)
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[AgentToolParams],
|
||||
) -> str:
|
||||
"""Execute the tool with the given parameters."""
|
||||
start_time = time.time()
|
||||
|
||||
# Create tool context
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
prompts = params.get("prompts")
|
||||
if prompts is None:
|
||||
await tool_ctx.error("No prompts provided")
|
||||
return "Error: At least one prompt must be provided."
|
||||
|
||||
# Handle both string and list inputs
|
||||
if isinstance(prompts, str):
|
||||
prompt_list = [prompts]
|
||||
elif isinstance(prompts, list):
|
||||
if not prompts:
|
||||
await tool_ctx.error("Empty prompts list provided")
|
||||
return "Error: At least one prompt must be provided."
|
||||
prompt_list = prompts
|
||||
else:
|
||||
await tool_ctx.error("Invalid prompts parameter type")
|
||||
return "Error: Parameter 'prompts' must be a string or list of strings."
|
||||
|
||||
# Validate absolute paths
|
||||
absolute_path_pattern = r"/(?:[^/\s]+/)*[^/\s]+"
|
||||
for prompt in prompt_list:
|
||||
if not re.search(absolute_path_pattern, prompt):
|
||||
await tool_ctx.error(f"Prompt missing absolute path: {prompt[:50]}...")
|
||||
return "Error: All prompts must contain at least one absolute path."
|
||||
|
||||
# Require hanzo-agents SDK
|
||||
if not HANZO_AGENTS_AVAILABLE:
|
||||
await tool_ctx.error("hanzo-agents SDK is required but not available")
|
||||
return "Error: hanzo-agents SDK is required for agent tool functionality. Please install it with: pip install hanzo-agents"
|
||||
|
||||
# Determine concurrency (parallel agents)
|
||||
concurrency = params.get("concurrency")
|
||||
if concurrency is not None and isinstance(concurrency, int) and concurrency > 0:
|
||||
# Expand prompt list to match concurrency
|
||||
if len(prompt_list) == 1:
|
||||
prompt_list = prompt_list * concurrency
|
||||
elif len(prompt_list) < concurrency:
|
||||
# Repeat prompts to reach concurrency
|
||||
times = (concurrency + len(prompt_list) - 1) // len(prompt_list)
|
||||
prompt_list = (prompt_list * times)[:concurrency]
|
||||
|
||||
await tool_ctx.info(
|
||||
f"Launching {len(prompt_list)} agent(s) using hanzo-agents SDK"
|
||||
)
|
||||
|
||||
# Determine model and agent type
|
||||
model = params.get("model", self.model_override)
|
||||
use_memory = params.get("use_memory", False)
|
||||
memory_backend = params.get("memory_backend", "sqlite")
|
||||
|
||||
# Get appropriate agent class
|
||||
agent_class = self._get_agent_class(model)
|
||||
|
||||
# Create state
|
||||
state = MCPAgentState(
|
||||
prompts=prompt_list,
|
||||
context={
|
||||
"permission_manager": self.permission_manager,
|
||||
"api_key": self.api_key_override,
|
||||
"base_url": self.base_url_override,
|
||||
"max_tokens": self.max_tokens_override,
|
||||
},
|
||||
)
|
||||
|
||||
# Create memory if requested
|
||||
memory_kv = None
|
||||
memory_vector = None
|
||||
if use_memory:
|
||||
memory_kv = create_memory_kv(memory_backend)
|
||||
memory_vector = create_memory_vector("simple")
|
||||
|
||||
# Create network
|
||||
network = Network(
|
||||
state=state,
|
||||
agents=[agent_class],
|
||||
router=sequential_router([agent_class] * len(prompt_list)),
|
||||
memory_kv=memory_kv,
|
||||
memory_vector=memory_vector,
|
||||
max_steps=self.max_iterations * len(prompt_list),
|
||||
)
|
||||
|
||||
# Execute
|
||||
try:
|
||||
final_state = await network.run()
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
# Format results
|
||||
results = final_state.results
|
||||
if len(results) == 1:
|
||||
formatted_result = f"""Agent execution completed in {execution_time:.2f} seconds.
|
||||
|
||||
AGENT RESPONSE:
|
||||
{results[0]}"""
|
||||
else:
|
||||
formatted_results = []
|
||||
for i, result in enumerate(results):
|
||||
formatted_results.append(f"Agent {i + 1} Result:\n{result}")
|
||||
|
||||
formatted_result = f"""Multi-agent execution completed in {execution_time:.2f} seconds ({len(results)} agents).
|
||||
|
||||
AGENT RESPONSES:
|
||||
{chr(10).join(formatted_results)}"""
|
||||
|
||||
await tool_ctx.info(f"Execution completed in {execution_time:.2f}s")
|
||||
return formatted_result
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Agent execution failed: {str(e)}")
|
||||
return f"Error: {str(e)}"
|
||||
|
||||
def _get_agent_class(self, model: Optional[str]) -> type[Agent]:
|
||||
"""Get appropriate agent class based on model."""
|
||||
if not model:
|
||||
model = "model://anthropic/claude-3-5-sonnet-20241022"
|
||||
|
||||
# Check for CLI agents
|
||||
cli_agents = {
|
||||
"claude_cli": ClaudeCodeAgent,
|
||||
"codex_cli": OpenAICodexAgent,
|
||||
"gemini_cli": GeminiAgent,
|
||||
"grok_cli": GrokAgent,
|
||||
}
|
||||
|
||||
if model in cli_agents:
|
||||
return cli_agents[model]
|
||||
|
||||
# Return generic MCP agent
|
||||
return type(
|
||||
"DynamicMCPAgent",
|
||||
(MCPAgent,),
|
||||
{
|
||||
"model": model,
|
||||
"__init__": lambda self: MCPAgent.__init__(
|
||||
self,
|
||||
available_tools=self.available_tools,
|
||||
permission_manager=self.permission_manager,
|
||||
ctx=self.ctx,
|
||||
model=model,
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this agent tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def dispatch_agent(
|
||||
prompts: str | list[str],
|
||||
ctx: MCPContext,
|
||||
model: Optional[str] = None,
|
||||
use_memory: bool = False,
|
||||
memory_backend: str = "sqlite",
|
||||
) -> str:
|
||||
return await tool_self.call(
|
||||
ctx,
|
||||
prompts=prompts,
|
||||
model=model,
|
||||
use_memory=use_memory,
|
||||
memory_backend=memory_backend,
|
||||
)
|
||||
@@ -1,645 +0,0 @@
|
||||
"""Agent tool implementation for Hanzo AI.
|
||||
|
||||
This module implements the AgentTool that allows Claude to delegate tasks to sub-agents,
|
||||
enabling concurrent execution of multiple operations and specialized processing.
|
||||
"""
|
||||
|
||||
import re
|
||||
import json
|
||||
import time
|
||||
import asyncio
|
||||
|
||||
# Import litellm with warnings suppressed
|
||||
import warnings
|
||||
from typing import Unpack, Annotated, TypedDict, final, override
|
||||
from collections.abc import Iterable
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", DeprecationWarning)
|
||||
import litellm
|
||||
from pydantic import Field
|
||||
from mcp.server import FastMCP
|
||||
from openai.types.chat import ChatCompletionToolParam, ChatCompletionMessageParam
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.jupyter import get_read_only_jupyter_tools
|
||||
from hanzo_mcp.tools.filesystem import Edit, MultiEdit, get_read_only_filesystem_tools
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.agent.prompt import (
|
||||
get_default_model,
|
||||
get_system_prompt,
|
||||
get_model_parameters,
|
||||
get_allowed_agent_tools,
|
||||
)
|
||||
from hanzo_mcp.tools.common.context import (
|
||||
ToolContext,
|
||||
create_tool_context,
|
||||
)
|
||||
from hanzo_mcp.tools.agent.critic_tool import CriticTool, CriticProtocol
|
||||
from hanzo_mcp.tools.agent.iching_tool import IChingTool
|
||||
from hanzo_mcp.tools.agent.review_tool import ReviewTool, ReviewProtocol
|
||||
from hanzo_mcp.tools.common.batch_tool import BatchTool
|
||||
from hanzo_mcp.tools.agent.tool_adapter import (
|
||||
convert_tools_to_openai_functions,
|
||||
)
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.agent.clarification_tool import ClarificationTool
|
||||
from hanzo_mcp.tools.agent.clarification_protocol import (
|
||||
ClarificationType,
|
||||
AgentClarificationMixin,
|
||||
)
|
||||
|
||||
Prompt = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Task for the agent to perform (must include absolute paths starting with /)",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class AgentToolParams(TypedDict, total=False):
|
||||
"""Parameters for the AgentTool.
|
||||
|
||||
Attributes:
|
||||
prompts: Task(s) for the agent to perform (must include absolute paths starting with /)
|
||||
"""
|
||||
|
||||
prompts: str | list[str]
|
||||
|
||||
|
||||
@final
|
||||
class AgentTool(AgentClarificationMixin, BaseTool):
|
||||
"""Tool for delegating tasks to sub-agents.
|
||||
|
||||
The AgentTool allows Claude to create and manage sub-agents for performing
|
||||
specialized tasks concurrently, such as code search, analysis, and more.
|
||||
|
||||
Agents can request clarification from the main loop up to once per task.
|
||||
"""
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name.
|
||||
|
||||
Returns:
|
||||
Tool name
|
||||
"""
|
||||
return "agent"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description.
|
||||
|
||||
Returns:
|
||||
Tool description
|
||||
"""
|
||||
# TODO: Add glob when it is implemented
|
||||
at = [t.name for t in self.available_tools]
|
||||
|
||||
return f"""Launch a new agent that has access to the following tools: {at}. When you are searching for a keyword or file and are not confident that you will find the right match in the first few tries, use the Agent tool to perform the search for you.
|
||||
|
||||
When to use the Agent tool:
|
||||
- If you are searching for a keyword like \"config\" or \"logger\", or for questions like \"which file does X?\", the Agent tool is strongly recommended
|
||||
- When you need to perform edits across multiple files based on search results
|
||||
- When you need to delegate complex file modification tasks
|
||||
|
||||
When NOT to use the Agent tool:
|
||||
- If you want to read a specific file path, use the read or glob tool instead of the Agent tool, to find the match more quickly
|
||||
- If you are searching for a specific class definition like \"class Foo\", use the glob tool instead, to find the match more quickly
|
||||
- If you are searching for code within a specific file or set of 2-3 files, use the read tool instead of the Agent tool, to find the match more quickly
|
||||
- Writing code and running bash commands (use other tools for that)
|
||||
- Other tasks that are not related to searching for a keyword or file
|
||||
|
||||
Usage notes:
|
||||
1. Launch multiple agents concurrently whenever possible, to maximize performance; to do that, use a single message with multiple tool uses
|
||||
2. When the agent is done, it will return a single message back to you. The result returned by the agent is not visible to the user. To show the user the result, you should send a text message back to the user with a concise summary of the result.
|
||||
3. Each agent invocation is stateless. You will not be able to send additional messages to the agent, nor will the agent be able to communicate with you outside of its final report. Therefore, your prompt should contain a highly detailed task description for the agent to perform autonomously and you should specify exactly what information the agent should return back to you in its final and only message to you.
|
||||
4. The agent's outputs should generally be trusted
|
||||
5. Clearly tell the agent whether you expect it to write code or just to do research (search, file reads, web fetches, etc.), since it is not aware of the user's intent"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
model: str | None = None,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
max_iterations: int = 10,
|
||||
max_tool_uses: int = 30,
|
||||
) -> None:
|
||||
"""Initialize the agent tool.
|
||||
|
||||
Args:
|
||||
|
||||
permission_manager: Permission manager for access control
|
||||
model: Optional model name override in LiteLLM format (e.g., "openai/gpt-4o")
|
||||
api_key: Optional API key for the model provider
|
||||
base_url: Optional base URL for the model provider API endpoint
|
||||
max_tokens: Optional maximum tokens for model responses
|
||||
max_iterations: Maximum number of iterations for agent (default: 10)
|
||||
max_tool_uses: Maximum number of total tool uses for agent (default: 30)
|
||||
"""
|
||||
|
||||
self.permission_manager = permission_manager
|
||||
self.model_override = model
|
||||
self.api_key_override = api_key
|
||||
self.base_url_override = base_url
|
||||
self.max_tokens_override = max_tokens
|
||||
self.max_iterations = max_iterations
|
||||
self.max_tool_uses = max_tool_uses
|
||||
self.available_tools: list[BaseTool] = []
|
||||
self.available_tools.extend(
|
||||
get_read_only_filesystem_tools(self.permission_manager)
|
||||
)
|
||||
self.available_tools.extend(
|
||||
get_read_only_jupyter_tools(self.permission_manager)
|
||||
)
|
||||
|
||||
# Always add edit tools - agents should have edit access
|
||||
self.available_tools.append(Edit(self.permission_manager))
|
||||
self.available_tools.append(MultiEdit(self.permission_manager))
|
||||
|
||||
# Add clarification tool for agents
|
||||
self.available_tools.append(ClarificationTool())
|
||||
|
||||
# Add critic tool for agents (devil's advocate)
|
||||
self.available_tools.append(CriticTool())
|
||||
|
||||
# Add review tool for agents (balanced review)
|
||||
self.available_tools.append(ReviewTool())
|
||||
|
||||
# Add I Ching tool for creative guidance
|
||||
self.available_tools.append(IChingTool())
|
||||
|
||||
self.available_tools.append(
|
||||
BatchTool({t.name: t for t in self.available_tools})
|
||||
)
|
||||
|
||||
# Initialize protocols
|
||||
self.critic_protocol = CriticProtocol()
|
||||
self.review_protocol = ReviewProtocol()
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[AgentToolParams],
|
||||
) -> str:
|
||||
"""Execute the tool with the given parameters.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Tool execution result
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Create tool context
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract and validate parameters
|
||||
prompts = params.get("prompts")
|
||||
|
||||
if prompts is None:
|
||||
await tool_ctx.error("No prompts provided")
|
||||
return """Error: At least one prompt must be provided.
|
||||
|
||||
IMPORTANT REMINDER FOR CLAUDE:
|
||||
The dispatch_agent tool requires prompts parameter. Please provide either:
|
||||
- A single prompt as a string
|
||||
- Multiple prompts as an array of strings
|
||||
|
||||
Each prompt must contain absolute paths starting with /.
|
||||
Example of correct usage:
|
||||
- prompts: "Search for all instances of the 'config' variable in /Users/bytedance/project/hanzo-mcp"
|
||||
- prompts: ["Find files in /path/to/project", "Search code in /path/to/src"]"""
|
||||
|
||||
# Handle both string and list inputs
|
||||
if isinstance(prompts, str):
|
||||
prompt_list = [prompts]
|
||||
elif isinstance(prompts, list):
|
||||
if not prompts:
|
||||
await tool_ctx.error("Empty prompts list provided")
|
||||
return "Error: At least one prompt must be provided when using a list."
|
||||
if not all(isinstance(p, str) for p in prompts):
|
||||
await tool_ctx.error("All prompts must be strings")
|
||||
return "Error: All prompts in the list must be strings."
|
||||
prompt_list = prompts
|
||||
else:
|
||||
await tool_ctx.error("Invalid prompts parameter type")
|
||||
return "Error: Parameter 'prompts' must be a string or an array of strings."
|
||||
|
||||
# Validate absolute paths in all prompts
|
||||
absolute_path_pattern = r"/(?:[^/\s]+/)*[^/\s]+"
|
||||
for prompt in prompt_list:
|
||||
if not re.search(absolute_path_pattern, prompt):
|
||||
await tool_ctx.error(
|
||||
f"Prompt does not contain absolute path: {prompt[:50]}..."
|
||||
)
|
||||
return """Error: All prompts must contain at least one absolute path.
|
||||
|
||||
IMPORTANT REMINDER FOR CLAUDE:
|
||||
When using the dispatch_agent tool, always include absolute paths in your prompts.
|
||||
Example of correct usage:
|
||||
- "Search for all instances of the 'config' variable in /Users/bytedance/project/hanzo-mcp"
|
||||
- "Find files that import the database module in /Users/bytedance/project/hanzo-mcp/src"
|
||||
|
||||
The agent cannot access files without knowing their absolute locations."""
|
||||
|
||||
# Execute agent(s) - always use _execute_multiple_agents for list inputs
|
||||
if len(prompt_list) == 1:
|
||||
await tool_ctx.info("Launching agent")
|
||||
result = await self._execute_multiple_agents(prompt_list, tool_ctx)
|
||||
execution_time = time.time() - start_time
|
||||
formatted_result = f"""Agent execution completed in {execution_time:.2f} seconds.
|
||||
|
||||
AGENT RESPONSE:
|
||||
{result}"""
|
||||
await tool_ctx.info(f"Agent execution completed in {execution_time:.2f}s")
|
||||
return formatted_result
|
||||
else:
|
||||
await tool_ctx.info(f"Launching {len(prompt_list)} agents in parallel")
|
||||
result = await self._execute_multiple_agents(prompt_list, tool_ctx)
|
||||
execution_time = time.time() - start_time
|
||||
formatted_result = f"""Multi-agent execution completed in {execution_time:.2f} seconds ({len(prompt_list)} agents).
|
||||
|
||||
AGENT RESPONSES:
|
||||
{result}"""
|
||||
await tool_ctx.info(
|
||||
f"Multi-agent execution completed in {execution_time:.2f}s"
|
||||
)
|
||||
return formatted_result
|
||||
|
||||
async def _execute_agent(self, prompt: str, tool_ctx: ToolContext) -> str:
|
||||
"""Execute a single agent with the given prompt.
|
||||
|
||||
Args:
|
||||
prompt: The task prompt for the agent
|
||||
tool_ctx: Tool context for logging
|
||||
|
||||
Returns:
|
||||
Agent execution result
|
||||
"""
|
||||
# Get available tools for the agent
|
||||
agent_tools = get_allowed_agent_tools(
|
||||
self.available_tools,
|
||||
self.permission_manager,
|
||||
)
|
||||
|
||||
# Convert tools to OpenAI format
|
||||
openai_tools = convert_tools_to_openai_functions(agent_tools)
|
||||
|
||||
# Log execution start
|
||||
await tool_ctx.info("Starting agent execution")
|
||||
|
||||
# Create a result container
|
||||
result = ""
|
||||
|
||||
try:
|
||||
# Create system prompt for this agent
|
||||
system_prompt = get_system_prompt(
|
||||
agent_tools,
|
||||
self.permission_manager,
|
||||
)
|
||||
|
||||
# Execute agent
|
||||
await tool_ctx.info(f"Executing agent task: {prompt[:50]}...")
|
||||
result = await self._execute_agent_with_tools(
|
||||
system_prompt, prompt, agent_tools, openai_tools, tool_ctx
|
||||
)
|
||||
except Exception as e:
|
||||
# Log and return error result
|
||||
error_message = f"Error executing agent: {str(e)}"
|
||||
await tool_ctx.error(error_message)
|
||||
return f"Error: {error_message}"
|
||||
|
||||
return result if result else "No results returned from agent"
|
||||
|
||||
async def _execute_multiple_agents(
|
||||
self, prompts: list[str], tool_ctx: ToolContext
|
||||
) -> str:
|
||||
"""Execute multiple agents concurrently.
|
||||
|
||||
Args:
|
||||
prompts: List of prompts for the agents
|
||||
tool_ctx: Tool context for logging
|
||||
|
||||
Returns:
|
||||
Combined results from all agents
|
||||
"""
|
||||
# Get available tools for the agents
|
||||
agent_tools = get_allowed_agent_tools(
|
||||
self.available_tools,
|
||||
self.permission_manager,
|
||||
)
|
||||
|
||||
# Convert tools to OpenAI format
|
||||
openai_tools = convert_tools_to_openai_functions(agent_tools)
|
||||
|
||||
# Create system prompt for the agents
|
||||
system_prompt = get_system_prompt(
|
||||
agent_tools,
|
||||
self.permission_manager,
|
||||
)
|
||||
|
||||
# Create tasks for parallel execution
|
||||
tasks = []
|
||||
for i, prompt in enumerate(prompts):
|
||||
await tool_ctx.info(f"Creating agent task {i + 1}: {prompt[:50]}...")
|
||||
task = self._execute_agent_with_tools(
|
||||
system_prompt, prompt, agent_tools, openai_tools, tool_ctx
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
# Execute all agents concurrently
|
||||
await tool_ctx.info(f"Executing {len(tasks)} agents in parallel")
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
# Handle single agent case
|
||||
if len(results) == 1:
|
||||
if isinstance(results[0], Exception):
|
||||
await tool_ctx.error(f"Agent execution failed: {str(results[0])}")
|
||||
return f"Error: {str(results[0])}"
|
||||
return results[0]
|
||||
|
||||
# Format results for multiple agents
|
||||
formatted_results = []
|
||||
for i, result in enumerate(results):
|
||||
if isinstance(result, Exception):
|
||||
formatted_results.append(f"Agent {i + 1} Error:\n{str(result)}")
|
||||
await tool_ctx.error(f"Agent {i + 1} failed: {str(result)}")
|
||||
else:
|
||||
formatted_results.append(f"Agent {i + 1} Result:\n{result}")
|
||||
|
||||
return "\n\n---\n\n".join(formatted_results)
|
||||
|
||||
async def _execute_agent_with_tools(
|
||||
self,
|
||||
system_prompt: str,
|
||||
user_prompt: str,
|
||||
available_tools: list[BaseTool],
|
||||
openai_tools: list[ChatCompletionToolParam],
|
||||
tool_ctx: ToolContext,
|
||||
) -> str:
|
||||
"""Execute agent with tool handling.
|
||||
|
||||
Args:
|
||||
system_prompt: System prompt for the agent
|
||||
user_prompt: User prompt for the agent
|
||||
available_tools: List of available tools
|
||||
openai_tools: List of tools in OpenAI format
|
||||
tool_ctx: Tool context for logging
|
||||
|
||||
Returns:
|
||||
Agent execution result
|
||||
"""
|
||||
# Get model parameters and name
|
||||
model = get_default_model(self.model_override)
|
||||
params = get_model_parameters(max_tokens=self.max_tokens_override)
|
||||
|
||||
# Initialize messages
|
||||
messages: Iterable[ChatCompletionMessageParam] = []
|
||||
messages.append({"role": "system", "content": system_prompt})
|
||||
messages.append({"role": "user", "content": user_prompt})
|
||||
|
||||
# Track tool usage for metrics
|
||||
tool_usage = {}
|
||||
total_tool_use_count = 0
|
||||
iteration_count = 0
|
||||
max_tool_uses = self.max_tool_uses # Safety limit to prevent infinite loops
|
||||
max_iterations = (
|
||||
self.max_iterations
|
||||
) # Add a maximum number of iterations for safety
|
||||
|
||||
# Execute until the agent completes or reaches the limit
|
||||
while total_tool_use_count < max_tool_uses and iteration_count < max_iterations:
|
||||
iteration_count += 1
|
||||
await tool_ctx.info(f"Calling model (iteration {iteration_count})...")
|
||||
|
||||
try:
|
||||
# Configure model parameters based on capabilities
|
||||
completion_params = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"tools": openai_tools,
|
||||
"tool_choice": "auto",
|
||||
"temperature": params["temperature"],
|
||||
"timeout": params["timeout"],
|
||||
}
|
||||
|
||||
if self.api_key_override:
|
||||
completion_params["api_key"] = self.api_key_override
|
||||
|
||||
# Add max_tokens if provided
|
||||
if params.get("max_tokens"):
|
||||
completion_params["max_tokens"] = params.get("max_tokens")
|
||||
|
||||
# Add base_url if provided
|
||||
if self.base_url_override:
|
||||
completion_params["base_url"] = self.base_url_override
|
||||
|
||||
# Make the model call
|
||||
response = litellm.completion(**completion_params) # pyright: ignore
|
||||
|
||||
if len(response.choices) == 0: # pyright: ignore
|
||||
raise ValueError("No response choices returned")
|
||||
|
||||
message = response.choices[0].message # pyright: ignore
|
||||
|
||||
# Add message to conversation history
|
||||
messages.append(message) # pyright: ignore
|
||||
|
||||
# If no tool calls, we're done
|
||||
if not message.tool_calls:
|
||||
return message.content or "Agent completed with no response."
|
||||
|
||||
# Process tool calls
|
||||
tool_call_count = len(message.tool_calls)
|
||||
await tool_ctx.info(f"Processing {tool_call_count} tool calls")
|
||||
|
||||
for tool_call in message.tool_calls:
|
||||
total_tool_use_count += 1
|
||||
function_name = tool_call.function.name
|
||||
|
||||
# Track usage
|
||||
tool_usage[function_name] = tool_usage.get(function_name, 0) + 1
|
||||
|
||||
# Log tool usage
|
||||
await tool_ctx.info(f"Agent using tool: {function_name}")
|
||||
|
||||
# Parse the arguments
|
||||
try:
|
||||
function_args = json.loads(tool_call.function.arguments)
|
||||
except json.JSONDecodeError:
|
||||
function_args = {}
|
||||
|
||||
# Find the matching tool
|
||||
tool = next(
|
||||
(t for t in available_tools if t.name == function_name), None
|
||||
)
|
||||
if not tool:
|
||||
tool_result = f"Error: Tool '{function_name}' not found"
|
||||
# Special handling for clarification requests
|
||||
elif function_name == "request_clarification":
|
||||
try:
|
||||
# Extract clarification parameters
|
||||
request_type = function_args.get("type", "ADDITIONAL_INFO")
|
||||
question = function_args.get("question", "")
|
||||
context = function_args.get("context", {})
|
||||
options = function_args.get("options")
|
||||
|
||||
# Convert string type to enum
|
||||
clarification_type = ClarificationType[request_type]
|
||||
|
||||
# Request clarification
|
||||
answer = await self.request_clarification(
|
||||
request_type=clarification_type,
|
||||
question=question,
|
||||
context=context,
|
||||
options=options,
|
||||
)
|
||||
|
||||
tool_result = self.format_clarification_in_output(
|
||||
question, answer
|
||||
)
|
||||
except Exception as e:
|
||||
tool_result = f"Error processing clarification: {str(e)}"
|
||||
# Special handling for critic requests
|
||||
elif function_name == "critic":
|
||||
try:
|
||||
# Extract critic parameters
|
||||
review_type = function_args.get("review_type", "GENERAL")
|
||||
work_description = function_args.get("work_description", "")
|
||||
code_snippets = function_args.get("code_snippets")
|
||||
file_paths = function_args.get("file_paths")
|
||||
specific_concerns = function_args.get("specific_concerns")
|
||||
|
||||
# Request critical review
|
||||
tool_result = self.critic_protocol.request_review(
|
||||
review_type=review_type,
|
||||
work_description=work_description,
|
||||
code_snippets=code_snippets,
|
||||
file_paths=file_paths,
|
||||
specific_concerns=specific_concerns,
|
||||
)
|
||||
except Exception as e:
|
||||
tool_result = f"Error processing critic review: {str(e)}"
|
||||
# Special handling for review requests
|
||||
elif function_name == "review":
|
||||
try:
|
||||
# Extract review parameters
|
||||
focus = function_args.get("focus", "GENERAL")
|
||||
work_description = function_args.get("work_description", "")
|
||||
code_snippets = function_args.get("code_snippets")
|
||||
file_paths = function_args.get("file_paths")
|
||||
context = function_args.get("context")
|
||||
|
||||
# Request balanced review
|
||||
tool_result = self.review_protocol.request_review(
|
||||
focus=focus,
|
||||
work_description=work_description,
|
||||
code_snippets=code_snippets,
|
||||
file_paths=file_paths,
|
||||
context=context,
|
||||
)
|
||||
except Exception as e:
|
||||
tool_result = f"Error processing review: {str(e)}"
|
||||
else:
|
||||
try:
|
||||
tool_result = await tool.call(
|
||||
ctx=tool_ctx.mcp_context, **function_args
|
||||
)
|
||||
except Exception as e:
|
||||
tool_result = f"Error executing {function_name}: {str(e)}"
|
||||
|
||||
await tool_ctx.info(
|
||||
f"tool {function_name} run with args {function_args} and return {tool_result[: min(100, len(tool_result))]}"
|
||||
)
|
||||
# Add the tool result to messages
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.id,
|
||||
"name": function_name,
|
||||
"content": tool_result,
|
||||
}
|
||||
)
|
||||
|
||||
# Log progress
|
||||
await tool_ctx.info(
|
||||
f"Processed {len(message.tool_calls)} tool calls. Total: {total_tool_use_count}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Error in model call: {str(e)}")
|
||||
# Avoid trying to JSON serialize message objects
|
||||
await tool_ctx.error(f"Message count: {len(messages)}")
|
||||
return f"Error in agent execution: {str(e)}"
|
||||
|
||||
# If we've reached the limit, add a warning and get final response
|
||||
if total_tool_use_count >= max_tool_uses or iteration_count >= max_iterations:
|
||||
messages.append(
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You have reached the maximum iteration. Please provide your final response.",
|
||||
}
|
||||
)
|
||||
|
||||
try:
|
||||
# Make a final call to get the result
|
||||
final_response = litellm.completion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=params["temperature"],
|
||||
timeout=params["timeout"],
|
||||
max_tokens=params.get("max_tokens"),
|
||||
)
|
||||
|
||||
return (
|
||||
final_response.choices[0].message.content
|
||||
or "Agent reached max iteration limit without a response."
|
||||
) # pyright: ignore
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Error in final model call: {str(e)}")
|
||||
return f"Error in final response: {str(e)}"
|
||||
|
||||
# Should not reach here but just in case
|
||||
return "Agent execution completed after maximum iterations."
|
||||
|
||||
def _format_result(self, result: str, execution_time: float) -> str:
|
||||
"""Format agent result with metrics.
|
||||
|
||||
Args:
|
||||
result: Raw result from agent
|
||||
execution_time: Execution time in seconds
|
||||
|
||||
Returns:
|
||||
Formatted result with metrics
|
||||
"""
|
||||
return f"""Agent execution completed in {execution_time:.2f} seconds.
|
||||
|
||||
AGENT RESPONSE:
|
||||
{result}
|
||||
"""
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this agent tool with the MCP server.
|
||||
|
||||
Creates a wrapper function with explicitly defined parameters that match
|
||||
the tool's parameter schema and registers it with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
"""
|
||||
tool_self = self # Create a reference to self for use in the closure
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def dispatch_agent(prompts: str | list[str], ctx: MCPContext) -> str:
|
||||
return await tool_self.call(ctx, prompts=prompts)
|
||||
@@ -1,225 +0,0 @@
|
||||
"""Clarification protocol for agent-to-mainloop communication.
|
||||
|
||||
This module provides a protocol for agents to request clarification
|
||||
from the main loop without human intervention.
|
||||
"""
|
||||
|
||||
import json
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
class ClarificationType(Enum):
|
||||
"""Types of clarification requests."""
|
||||
|
||||
AMBIGUOUS_INSTRUCTION = "ambiguous_instruction"
|
||||
MISSING_CONTEXT = "missing_context"
|
||||
MULTIPLE_OPTIONS = "multiple_options"
|
||||
CONFIRMATION_NEEDED = "confirmation_needed"
|
||||
ADDITIONAL_INFO = "additional_info"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClarificationRequest:
|
||||
"""A request for clarification from an agent."""
|
||||
|
||||
agent_id: str
|
||||
request_type: ClarificationType
|
||||
question: str
|
||||
context: Dict[str, Any]
|
||||
options: Optional[List[str]] = None
|
||||
|
||||
def to_json(self) -> str:
|
||||
"""Convert to JSON for transport."""
|
||||
return json.dumps(
|
||||
{
|
||||
"agent_id": self.agent_id,
|
||||
"request_type": self.request_type.value,
|
||||
"question": self.question,
|
||||
"context": self.context,
|
||||
"options": self.options,
|
||||
}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, data: str) -> "ClarificationRequest":
|
||||
"""Create from JSON string."""
|
||||
obj = json.loads(data)
|
||||
return cls(
|
||||
agent_id=obj["agent_id"],
|
||||
request_type=ClarificationType(obj["request_type"]),
|
||||
question=obj["question"],
|
||||
context=obj["context"],
|
||||
options=obj.get("options"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClarificationResponse:
|
||||
"""A response to a clarification request."""
|
||||
|
||||
request_id: str
|
||||
answer: str
|
||||
additional_context: Optional[Dict[str, Any]] = None
|
||||
|
||||
def to_json(self) -> str:
|
||||
"""Convert to JSON for transport."""
|
||||
return json.dumps(
|
||||
{
|
||||
"request_id": self.request_id,
|
||||
"answer": self.answer,
|
||||
"additional_context": self.additional_context,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class ClarificationHandler:
|
||||
"""Handles clarification requests from agents."""
|
||||
|
||||
def __init__(self):
|
||||
self.pending_requests: Dict[str, ClarificationRequest] = {}
|
||||
self.request_counter = 0
|
||||
|
||||
def create_request(
|
||||
self,
|
||||
agent_id: str,
|
||||
request_type: ClarificationType,
|
||||
question: str,
|
||||
context: Dict[str, Any],
|
||||
options: Optional[List[str]] = None,
|
||||
) -> str:
|
||||
"""Create a new clarification request.
|
||||
|
||||
Returns:
|
||||
Request ID for tracking
|
||||
"""
|
||||
request = ClarificationRequest(
|
||||
agent_id=agent_id,
|
||||
request_type=request_type,
|
||||
question=question,
|
||||
context=context,
|
||||
options=options,
|
||||
)
|
||||
|
||||
request_id = f"clarify_{self.request_counter}"
|
||||
self.request_counter += 1
|
||||
self.pending_requests[request_id] = request
|
||||
|
||||
return request_id
|
||||
|
||||
def handle_request(self, request: ClarificationRequest) -> ClarificationResponse:
|
||||
"""Handle a clarification request automatically.
|
||||
|
||||
This method implements automatic clarification resolution
|
||||
based on context and common patterns.
|
||||
"""
|
||||
request_id = f"clarify_{len(self.pending_requests)}"
|
||||
|
||||
# Handle different types of clarification
|
||||
if request.request_type == ClarificationType.AMBIGUOUS_INSTRUCTION:
|
||||
# Try to clarify based on context
|
||||
if "file_path" in request.context:
|
||||
if request.context["file_path"].endswith(".go"):
|
||||
answer = "For Go files, ensure you add imports in the correct format and handle both single import and import block cases."
|
||||
elif request.context["file_path"].endswith(".py"):
|
||||
answer = "For Python files, add imports at the top of the file after any module docstring."
|
||||
else:
|
||||
answer = "Add imports according to the language's conventions."
|
||||
else:
|
||||
answer = "Proceed with the most reasonable interpretation based on the context."
|
||||
|
||||
elif request.request_type == ClarificationType.MISSING_CONTEXT:
|
||||
# Provide additional context based on what's missing
|
||||
if "import_path" in request.question.lower():
|
||||
answer = "Use the standard import path based on the project structure. Check existing imports in similar files for patterns."
|
||||
elif "format" in request.question.lower():
|
||||
answer = "Match the existing code style in the file. Use the same indentation and formatting patterns."
|
||||
else:
|
||||
answer = "Analyze the surrounding code and project structure to infer the missing information."
|
||||
|
||||
elif request.request_type == ClarificationType.MULTIPLE_OPTIONS:
|
||||
# Choose the best option based on context
|
||||
if request.options:
|
||||
# Simple heuristic: choose the first option that seems most standard
|
||||
for option in request.options:
|
||||
if "common" in option or "standard" in option:
|
||||
answer = f"Choose option: {option}"
|
||||
break
|
||||
else:
|
||||
answer = f"Choose option: {request.options[0]}"
|
||||
else:
|
||||
answer = "Choose the most conventional approach based on the codebase patterns."
|
||||
|
||||
elif request.request_type == ClarificationType.CONFIRMATION_NEEDED:
|
||||
# Auto-confirm safe operations
|
||||
if "add import" in request.question.lower():
|
||||
answer = "Yes, proceed with adding the import."
|
||||
elif "multi_edit" in request.question.lower():
|
||||
answer = "Yes, use multi_edit for efficiency."
|
||||
else:
|
||||
answer = "Proceed if the operation is safe and reversible."
|
||||
|
||||
else: # ADDITIONAL_INFO
|
||||
answer = "Continue with available information and make reasonable assumptions based on context."
|
||||
|
||||
return ClarificationResponse(
|
||||
request_id=request_id,
|
||||
answer=answer,
|
||||
additional_context={"auto_resolved": True},
|
||||
)
|
||||
|
||||
|
||||
class AgentClarificationMixin:
|
||||
"""Mixin for agents to request clarification."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.clarification_handler = ClarificationHandler()
|
||||
self.clarification_count = 0
|
||||
self.max_clarifications = 1 # Allow up to 1 clarification per task
|
||||
|
||||
async def request_clarification(
|
||||
self,
|
||||
request_type: ClarificationType,
|
||||
question: str,
|
||||
context: Dict[str, Any],
|
||||
options: Optional[List[str]] = None,
|
||||
) -> str:
|
||||
"""Request clarification from the main loop.
|
||||
|
||||
Args:
|
||||
request_type: Type of clarification needed
|
||||
question: The question to ask
|
||||
context: Relevant context for the question
|
||||
options: Optional list of choices
|
||||
|
||||
Returns:
|
||||
The clarification response
|
||||
|
||||
Raises:
|
||||
RuntimeError: If clarification limit exceeded
|
||||
"""
|
||||
if self.clarification_count >= self.max_clarifications:
|
||||
raise RuntimeError("Clarification limit exceeded")
|
||||
|
||||
self.clarification_count += 1
|
||||
|
||||
# Create request
|
||||
request = ClarificationRequest(
|
||||
agent_id=getattr(self, "agent_id", "unknown"),
|
||||
request_type=request_type,
|
||||
question=question,
|
||||
context=context,
|
||||
options=options,
|
||||
)
|
||||
|
||||
# In real implementation, this would communicate with main loop
|
||||
# For now, use the automatic handler
|
||||
response = self.clarification_handler.handle_request(request)
|
||||
|
||||
return response.answer
|
||||
|
||||
def format_clarification_in_output(self, question: str, answer: str) -> str:
|
||||
"""Format clarification exchange for output."""
|
||||
return f"\n🤔 Clarification needed: {question}\n✅ Resolved: {answer}\n"
|
||||
@@ -1,69 +0,0 @@
|
||||
"""Clarification tool for agents to request information from main loop."""
|
||||
|
||||
from typing import Any, Dict, List, Optional, override
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
|
||||
|
||||
class ClarificationTool(BaseTool):
|
||||
"""Tool for agents to request clarification from the main loop."""
|
||||
|
||||
name = "request_clarification"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Request clarification from the main loop (not the human user).
|
||||
|
||||
Use this when you encounter:
|
||||
- Ambiguous instructions that could be interpreted multiple ways
|
||||
- Missing context needed to complete the task
|
||||
- Multiple valid options where you need guidance
|
||||
- Operations that need confirmation before proceeding
|
||||
- Need for additional information not provided
|
||||
|
||||
Parameters:
|
||||
- type: Type of clarification (AMBIGUOUS_INSTRUCTION, MISSING_CONTEXT, MULTIPLE_OPTIONS, CONFIRMATION_NEEDED, ADDITIONAL_INFO)
|
||||
- question: Clear, specific question to ask
|
||||
- context: Relevant context (e.g., file_path, current_operation, etc.)
|
||||
- options: Optional list of possible choices (for MULTIPLE_OPTIONS type)
|
||||
|
||||
You can only use this ONCE per task, so make it count!
|
||||
|
||||
Example:
|
||||
request_clarification(
|
||||
type="MISSING_CONTEXT",
|
||||
question="What is the correct import path for the common package?",
|
||||
context={"file_path": "/path/to/file.go", "undefined_symbol": "common"},
|
||||
options=["github.com/luxfi/node/common", "github.com/project/common"]
|
||||
)"""
|
||||
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
type: str,
|
||||
question: str,
|
||||
context: Dict[str, Any],
|
||||
options: Optional[List[str]] = None,
|
||||
) -> str:
|
||||
"""This is a placeholder - actual implementation happens in AgentTool."""
|
||||
# This tool is handled specially in the agent execution
|
||||
return f"Clarification request: {question}"
|
||||
|
||||
def register(self, server: FastMCP) -> None:
|
||||
"""Register the tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@server.tool(name=self.name, description=self.description)
|
||||
async def request_clarification(
|
||||
ctx: MCPContext,
|
||||
type: str,
|
||||
question: str,
|
||||
context: Dict[str, Any],
|
||||
options: Optional[List[str]] = None,
|
||||
) -> str:
|
||||
return await tool_self.call(ctx, type, question, context, options)
|
||||
@@ -1,126 +0,0 @@
|
||||
"""Claude Code CLI agent tool.
|
||||
|
||||
This tool provides integration with the Claude Code CLI (claude command),
|
||||
allowing programmatic execution of Claude for code tasks.
|
||||
"""
|
||||
|
||||
from typing import List, Optional, final, override
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.agent.code_auth import get_latest_claude_model
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.agent.cli_agent_base import CLIAgentBase
|
||||
|
||||
|
||||
@final
|
||||
class ClaudeCLITool(CLIAgentBase):
|
||||
"""Tool for executing Claude Code CLI."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
model: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize Claude CLI tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
model: Optional model override (defaults to latest Sonnet)
|
||||
**kwargs: Additional arguments
|
||||
"""
|
||||
super().__init__(
|
||||
permission_manager=permission_manager,
|
||||
command_name="claude",
|
||||
provider_name="Claude Code",
|
||||
default_model=model or get_latest_claude_model(),
|
||||
env_vars=["ANTHROPIC_API_KEY", "CLAUDE_API_KEY"],
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "claude_cli"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Execute Claude Code CLI for advanced code tasks.
|
||||
|
||||
This tool runs the Claude Code CLI (claude command) for code generation,
|
||||
editing, analysis, and other programming tasks. It uses the latest
|
||||
Claude 3.5 Sonnet model by default.
|
||||
|
||||
Features:
|
||||
- Direct access to Claude's coding capabilities
|
||||
- File-aware context and editing
|
||||
- Interactive code generation
|
||||
- Supports all Claude Code CLI features
|
||||
|
||||
Usage:
|
||||
claude_cli(prompts="Fix the bug in main.py and add tests")
|
||||
claude_cli(prompts="Refactor this class to use dependency injection", model="claude-3-opus-20240229")
|
||||
|
||||
Requirements:
|
||||
- Claude Code CLI must be installed
|
||||
- ANTHROPIC_API_KEY or CLAUDE_API_KEY environment variable
|
||||
"""
|
||||
|
||||
@override
|
||||
def get_cli_args(self, prompt: str, **kwargs) -> List[str]:
|
||||
"""Get CLI arguments for Claude.
|
||||
|
||||
Args:
|
||||
prompt: The prompt to send
|
||||
**kwargs: Additional arguments (model, temperature, etc.)
|
||||
|
||||
Returns:
|
||||
List of command arguments
|
||||
"""
|
||||
args = []
|
||||
|
||||
# Add model if specified
|
||||
model = kwargs.get("model", self.default_model)
|
||||
if model:
|
||||
args.extend(["--model", model])
|
||||
|
||||
# Add temperature if specified
|
||||
if "temperature" in kwargs:
|
||||
args.extend(["--temperature", str(kwargs["temperature"])])
|
||||
|
||||
# Add max tokens if specified
|
||||
if "max_tokens" in kwargs:
|
||||
args.extend(["--max-tokens", str(kwargs["max_tokens"])])
|
||||
|
||||
# Add the prompt
|
||||
args.append(prompt)
|
||||
|
||||
return args
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def claude_cli(
|
||||
ctx: MCPContext,
|
||||
prompts: str,
|
||||
model: Optional[str] = None,
|
||||
temperature: Optional[float] = None,
|
||||
max_tokens: Optional[int] = None,
|
||||
working_dir: Optional[str] = None,
|
||||
) -> str:
|
||||
return await tool_self.call(
|
||||
ctx,
|
||||
prompts=prompts,
|
||||
model=model,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
working_dir=working_dir,
|
||||
)
|
||||
@@ -1,494 +0,0 @@
|
||||
"""Claude Desktop authentication management.
|
||||
|
||||
This module provides tools to automate Claude Desktop login/logout,
|
||||
manage separate accounts for swarm agents, and handle authentication flows.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
import asyncio
|
||||
import subprocess
|
||||
import webbrowser
|
||||
from typing import Any, Dict, Tuple, Optional
|
||||
from pathlib import Path
|
||||
from urllib.parse import parse_qs
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
|
||||
|
||||
class ClaudeDesktopAuth:
|
||||
"""Manages Claude Desktop authentication."""
|
||||
|
||||
# Claude Desktop paths
|
||||
CLAUDE_APP_MAC = "/Applications/Claude.app"
|
||||
CLAUDE_CONFIG_DIR = Path.home() / ".claude"
|
||||
CLAUDE_SESSION_FILE = CLAUDE_CONFIG_DIR / "session.json"
|
||||
CLAUDE_ACCOUNTS_FILE = CLAUDE_CONFIG_DIR / "accounts.json"
|
||||
|
||||
# Authentication endpoints
|
||||
CLAUDE_LOGIN_URL = "https://claude.ai/login"
|
||||
CLAUDE_API_URL = "https://api.claude.ai"
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize Claude Desktop auth manager."""
|
||||
self.ensure_config_dir()
|
||||
|
||||
def ensure_config_dir(self):
|
||||
"""Ensure Claude config directory exists."""
|
||||
self.CLAUDE_CONFIG_DIR.mkdir(exist_ok=True)
|
||||
|
||||
def is_claude_installed(self) -> bool:
|
||||
"""Check if Claude Desktop is installed."""
|
||||
if os.path.exists(self.CLAUDE_APP_MAC):
|
||||
return True
|
||||
|
||||
# Check if claude command is available
|
||||
try:
|
||||
result = subprocess.run(["which", "claude"], capture_output=True, text=True)
|
||||
return result.returncode == 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def is_logged_in(self, account: Optional[str] = None) -> bool:
|
||||
"""Check if Claude Desktop is logged in.
|
||||
|
||||
Args:
|
||||
account: Optional account identifier to check
|
||||
|
||||
Returns:
|
||||
True if logged in
|
||||
"""
|
||||
if not self.CLAUDE_SESSION_FILE.exists():
|
||||
return False
|
||||
|
||||
try:
|
||||
with open(self.CLAUDE_SESSION_FILE, "r") as f:
|
||||
session = json.load(f)
|
||||
|
||||
# Check if session is valid
|
||||
if not session.get("access_token"):
|
||||
return False
|
||||
|
||||
# Check expiry if available
|
||||
if "expires_at" in session:
|
||||
if time.time() > session["expires_at"]:
|
||||
return False
|
||||
|
||||
# Check specific account if requested
|
||||
if account and session.get("account") != account:
|
||||
return False
|
||||
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def get_current_account(self) -> Optional[str]:
|
||||
"""Get the currently logged in account."""
|
||||
if not self.is_logged_in():
|
||||
return None
|
||||
|
||||
try:
|
||||
with open(self.CLAUDE_SESSION_FILE, "r") as f:
|
||||
session = json.load(f)
|
||||
return session.get("account", session.get("email"))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
async def login_interactive(
|
||||
self, account: Optional[str] = None, headless: bool = False
|
||||
) -> Tuple[bool, str]:
|
||||
"""Login to Claude Desktop interactively.
|
||||
|
||||
Args:
|
||||
account: Optional account email/identifier
|
||||
headless: Whether to run in headless mode
|
||||
|
||||
Returns:
|
||||
Tuple of (success, message)
|
||||
"""
|
||||
# Check if already logged in
|
||||
if self.is_logged_in(account):
|
||||
current = self.get_current_account()
|
||||
return True, f"Already logged in as {current}"
|
||||
|
||||
# Start login flow
|
||||
if headless:
|
||||
return await self._login_headless(account)
|
||||
else:
|
||||
return await self._login_browser(account)
|
||||
|
||||
async def _login_browser(self, account: Optional[str]) -> Tuple[bool, str]:
|
||||
"""Login using browser flow."""
|
||||
# Generate state for OAuth-like flow
|
||||
state = os.urandom(16).hex()
|
||||
|
||||
# Create callback server
|
||||
callback_port = 9876
|
||||
auth_code = None
|
||||
|
||||
async def handle_callback(reader, writer):
|
||||
"""Handle OAuth callback."""
|
||||
nonlocal auth_code
|
||||
|
||||
# Read request
|
||||
request = await reader.read(1024)
|
||||
request_str = request.decode()
|
||||
|
||||
# Extract code from query params
|
||||
if "GET /" in request_str:
|
||||
path = request_str.split(" ")[1]
|
||||
if "?code=" in path:
|
||||
query = path.split("?")[1]
|
||||
params = parse_qs(query)
|
||||
if "code" in params:
|
||||
auth_code = params["code"][0]
|
||||
|
||||
# Send response
|
||||
response = b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n"
|
||||
response += b"<html><body><h1>Authentication successful!</h1>"
|
||||
response += b"<p>You can close this window.</p></body></html>"
|
||||
writer.write(response)
|
||||
await writer.drain()
|
||||
writer.close()
|
||||
|
||||
# Start callback server
|
||||
server = await asyncio.start_server(handle_callback, "localhost", callback_port)
|
||||
|
||||
# Build login URL
|
||||
login_url = f"{self.CLAUDE_LOGIN_URL}?callback=http://localhost:{callback_port}&state={state}"
|
||||
if account:
|
||||
login_url += f"&login_hint={account}"
|
||||
|
||||
# Open browser
|
||||
print(f"Opening browser for Claude login...")
|
||||
print(f"URL: {login_url}")
|
||||
webbrowser.open(login_url)
|
||||
|
||||
# Wait for callback (timeout after 2 minutes)
|
||||
try:
|
||||
start_time = time.time()
|
||||
while not auth_code and (time.time() - start_time) < 120:
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
if auth_code:
|
||||
# Exchange code for session
|
||||
success = await self._exchange_code_for_session(auth_code, account)
|
||||
if success:
|
||||
return True, f"Successfully logged in as {account or 'default'}"
|
||||
else:
|
||||
return False, "Failed to exchange auth code for session"
|
||||
else:
|
||||
return False, "Login timeout - no auth code received"
|
||||
|
||||
finally:
|
||||
server.close()
|
||||
await server.wait_closed()
|
||||
|
||||
async def _login_headless(self, account: Optional[str]) -> Tuple[bool, str]:
|
||||
"""Login in headless mode using TTY automation."""
|
||||
# This would use expect/pexpect or similar to automate the CLI
|
||||
# For now, return a placeholder
|
||||
return False, "Headless login not yet implemented"
|
||||
|
||||
async def _exchange_code_for_session(
|
||||
self, code: str, account: Optional[str]
|
||||
) -> bool:
|
||||
"""Exchange auth code for session token."""
|
||||
# This would make API calls to exchange the code
|
||||
# For now, create a mock session
|
||||
session = {
|
||||
"access_token": f"mock_token_{code[:8]}",
|
||||
"account": account or "default",
|
||||
"email": account,
|
||||
"expires_at": time.time() + 3600 * 24, # 24 hours
|
||||
"created_at": time.time(),
|
||||
}
|
||||
|
||||
try:
|
||||
with open(self.CLAUDE_SESSION_FILE, "w") as f:
|
||||
json.dump(session, f, indent=2)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def logout(self, account: Optional[str] = None) -> Tuple[bool, str]:
|
||||
"""Logout from Claude Desktop.
|
||||
|
||||
Args:
|
||||
account: Optional account to logout (if multiple accounts)
|
||||
|
||||
Returns:
|
||||
Tuple of (success, message)
|
||||
"""
|
||||
current = self.get_current_account()
|
||||
|
||||
if not current:
|
||||
return True, "No active session to logout"
|
||||
|
||||
if account and current != account:
|
||||
return False, f"Not logged in as {account} (current: {current})"
|
||||
|
||||
try:
|
||||
# Remove session file
|
||||
if self.CLAUDE_SESSION_FILE.exists():
|
||||
self.CLAUDE_SESSION_FILE.unlink()
|
||||
|
||||
# Clear any cached credentials
|
||||
self._clear_credentials_cache()
|
||||
|
||||
return True, f"Successfully logged out {current}"
|
||||
except Exception as e:
|
||||
return False, f"Logout failed: {str(e)}"
|
||||
|
||||
def _clear_credentials_cache(self):
|
||||
"""Clear any cached credentials."""
|
||||
# Clear keychain on macOS
|
||||
if os.path.exists("/usr/bin/security"):
|
||||
try:
|
||||
subprocess.run(
|
||||
[
|
||||
"/usr/bin/security",
|
||||
"delete-generic-password",
|
||||
"-s",
|
||||
"claude.ai",
|
||||
"-a",
|
||||
"claude-desktop",
|
||||
],
|
||||
capture_output=True,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def switch_account(self, account: str) -> Tuple[bool, str]:
|
||||
"""Switch to a different Claude account.
|
||||
|
||||
Args:
|
||||
account: Account identifier to switch to
|
||||
|
||||
Returns:
|
||||
Tuple of (success, message)
|
||||
"""
|
||||
# Load accounts configuration
|
||||
accounts = self._load_accounts()
|
||||
|
||||
if account not in accounts:
|
||||
return False, f"Unknown account: {account}"
|
||||
|
||||
# Save current session if any
|
||||
current = self.get_current_account()
|
||||
if current and current != account:
|
||||
self._save_session_for_account(current)
|
||||
|
||||
# Load session for new account
|
||||
if self._load_session_for_account(account):
|
||||
return True, f"Switched to account: {account}"
|
||||
else:
|
||||
return False, f"No saved session for account: {account}"
|
||||
|
||||
def _load_accounts(self) -> Dict[str, Any]:
|
||||
"""Load accounts configuration."""
|
||||
if not self.CLAUDE_ACCOUNTS_FILE.exists():
|
||||
return {}
|
||||
|
||||
try:
|
||||
with open(self.CLAUDE_ACCOUNTS_FILE, "r") as f:
|
||||
return json.load(f)
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
def _save_accounts(self, accounts: Dict[str, Any]):
|
||||
"""Save accounts configuration."""
|
||||
with open(self.CLAUDE_ACCOUNTS_FILE, "w") as f:
|
||||
json.dump(accounts, f, indent=2)
|
||||
|
||||
def _save_session_for_account(self, account: str):
|
||||
"""Save current session for an account."""
|
||||
if not self.CLAUDE_SESSION_FILE.exists():
|
||||
return
|
||||
|
||||
accounts = self._load_accounts()
|
||||
|
||||
try:
|
||||
with open(self.CLAUDE_SESSION_FILE, "r") as f:
|
||||
session = json.load(f)
|
||||
|
||||
accounts[account] = {"session": session, "saved_at": time.time()}
|
||||
|
||||
self._save_accounts(accounts)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _load_session_for_account(self, account: str) -> bool:
|
||||
"""Load saved session for an account."""
|
||||
accounts = self._load_accounts()
|
||||
|
||||
if account not in accounts:
|
||||
return False
|
||||
|
||||
account_data = accounts[account]
|
||||
if "session" not in account_data:
|
||||
return False
|
||||
|
||||
try:
|
||||
# Restore session
|
||||
session = account_data["session"]
|
||||
|
||||
# Update account info
|
||||
session["account"] = account
|
||||
|
||||
with open(self.CLAUDE_SESSION_FILE, "w") as f:
|
||||
json.dump(session, f, indent=2)
|
||||
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def create_agent_account(self, agent_id: str) -> str:
|
||||
"""Create a unique account identifier for an agent.
|
||||
|
||||
Args:
|
||||
agent_id: Unique agent identifier
|
||||
|
||||
Returns:
|
||||
Account identifier for the agent
|
||||
"""
|
||||
# Generate agent-specific account
|
||||
return f"agent_{agent_id}@claude.local"
|
||||
|
||||
async def ensure_agent_auth(
|
||||
self, agent_id: str, force_new: bool = False
|
||||
) -> Tuple[bool, str]:
|
||||
"""Ensure an agent is authenticated with its own account.
|
||||
|
||||
Args:
|
||||
agent_id: Unique agent identifier
|
||||
force_new: Force new login even if cached
|
||||
|
||||
Returns:
|
||||
Tuple of (success, message/account)
|
||||
"""
|
||||
agent_account = self.create_agent_account(agent_id)
|
||||
|
||||
# Check if agent already has a session
|
||||
if not force_new and self._has_saved_session(agent_account):
|
||||
# Try to switch to agent account
|
||||
success, msg = self.switch_account(agent_account)
|
||||
if success:
|
||||
return True, agent_account
|
||||
|
||||
# Need to create new session for agent
|
||||
# For now, we'll use the main account
|
||||
# In production, this would create separate auth
|
||||
current = self.get_current_account()
|
||||
if current:
|
||||
# Clone current session for agent
|
||||
self._clone_session_for_agent(current, agent_account)
|
||||
return True, agent_account
|
||||
else:
|
||||
return False, "No active session to clone for agent"
|
||||
|
||||
def _has_saved_session(self, account: str) -> bool:
|
||||
"""Check if account has a saved session."""
|
||||
accounts = self._load_accounts()
|
||||
return account in accounts and "session" in accounts[account]
|
||||
|
||||
def _clone_session_for_agent(self, source: str, agent_account: str):
|
||||
"""Clone a session for an agent account."""
|
||||
# In a real implementation, this would create a sub-session
|
||||
# or use delegation tokens
|
||||
if self.CLAUDE_SESSION_FILE.exists():
|
||||
try:
|
||||
with open(self.CLAUDE_SESSION_FILE, "r") as f:
|
||||
session = json.load(f)
|
||||
|
||||
# Modify for agent
|
||||
session["account"] = agent_account
|
||||
session["parent_account"] = source
|
||||
session["is_agent"] = True
|
||||
|
||||
# Save as agent session
|
||||
accounts = self._load_accounts()
|
||||
accounts[agent_account] = {
|
||||
"session": session,
|
||||
"saved_at": time.time(),
|
||||
"parent": source,
|
||||
}
|
||||
self._save_accounts(accounts)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class ClaudeDesktopAuthTool(BaseTool):
|
||||
"""Tool for managing Claude Desktop authentication."""
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "claude_auth"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return """Manage Claude Desktop authentication.
|
||||
|
||||
Actions:
|
||||
- status: Check login status
|
||||
- login: Login to Claude Desktop
|
||||
- logout: Logout from Claude Desktop
|
||||
- switch: Switch between accounts
|
||||
- ensure_agent: Ensure agent has auth
|
||||
|
||||
Usage:
|
||||
claude_auth status
|
||||
claude_auth login --account [email protected]
|
||||
claude_auth logout
|
||||
claude_auth switch agent_1
|
||||
claude_auth ensure_agent swarm_agent_1"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the auth tool."""
|
||||
self.auth = ClaudeDesktopAuth()
|
||||
|
||||
async def call(self, ctx, action: str = "status", **kwargs) -> str:
|
||||
"""Execute auth action."""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
if action == "status":
|
||||
if self.auth.is_logged_in():
|
||||
account = self.auth.get_current_account()
|
||||
return f"Logged in as: {account}"
|
||||
else:
|
||||
return "Not logged in"
|
||||
|
||||
elif action == "login":
|
||||
account = kwargs.get("account")
|
||||
headless = kwargs.get("headless", False)
|
||||
success, msg = await self.auth.login_interactive(account, headless)
|
||||
return msg
|
||||
|
||||
elif action == "logout":
|
||||
account = kwargs.get("account")
|
||||
success, msg = await self.auth.logout(account)
|
||||
return msg
|
||||
|
||||
elif action == "switch":
|
||||
account = kwargs.get("account")
|
||||
if not account:
|
||||
return "Error: account required for switch"
|
||||
success, msg = self.auth.switch_account(account)
|
||||
return msg
|
||||
|
||||
elif action == "ensure_agent":
|
||||
agent_id = kwargs.get("agent_id")
|
||||
if not agent_id:
|
||||
return "Error: agent_id required"
|
||||
force_new = kwargs.get("force_new", False)
|
||||
success, result = await self.auth.ensure_agent_auth(agent_id, force_new)
|
||||
if success:
|
||||
return f"Agent authenticated as: {result}"
|
||||
else:
|
||||
return f"Failed: {result}"
|
||||
|
||||
else:
|
||||
return f"Unknown action: {action}"
|
||||
@@ -1,197 +0,0 @@
|
||||
"""Base class for CLI-based AI agent tools.
|
||||
|
||||
This provides common functionality for spawning CLI-based AI coding assistants
|
||||
like Claude Code, OpenAI Codex, Google Gemini, and Grok.
|
||||
"""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import asyncio
|
||||
import tempfile
|
||||
from abc import abstractmethod
|
||||
from typing import List, Optional
|
||||
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
|
||||
class CLIAgentBase(BaseTool):
|
||||
"""Base class for CLI-based AI agent tools."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
command_name: str,
|
||||
provider_name: str,
|
||||
default_model: Optional[str] = None,
|
||||
env_vars: Optional[List[str]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize CLI agent base.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
command_name: The CLI command name (e.g., 'claude', 'openai')
|
||||
provider_name: The provider name (e.g., 'Claude', 'OpenAI')
|
||||
default_model: Default model to use
|
||||
env_vars: List of environment variables to check for API keys
|
||||
**kwargs: Additional arguments
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.command_name = command_name
|
||||
self.provider_name = provider_name
|
||||
self.default_model = default_model
|
||||
self.env_vars = env_vars or []
|
||||
|
||||
def is_installed(self) -> bool:
|
||||
"""Check if the CLI tool is installed."""
|
||||
return shutil.which(self.command_name) is not None
|
||||
|
||||
def has_api_key(self) -> bool:
|
||||
"""Check if API key is available in environment."""
|
||||
if not self.env_vars:
|
||||
return True # No API key needed
|
||||
|
||||
for var in self.env_vars:
|
||||
if os.environ.get(var):
|
||||
return True
|
||||
return False
|
||||
|
||||
@abstractmethod
|
||||
def get_cli_args(self, prompt: str, **kwargs) -> List[str]:
|
||||
"""Get CLI arguments for the specific tool.
|
||||
|
||||
Args:
|
||||
prompt: The prompt to send
|
||||
**kwargs: Additional arguments
|
||||
|
||||
Returns:
|
||||
List of command arguments
|
||||
"""
|
||||
pass
|
||||
|
||||
async def execute_cli(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
prompt: str,
|
||||
working_dir: Optional[str] = None,
|
||||
timeout: int = 300,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
"""Execute the CLI command.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
prompt: The prompt to send
|
||||
working_dir: Working directory for the command
|
||||
timeout: Command timeout in seconds
|
||||
**kwargs: Additional arguments
|
||||
|
||||
Returns:
|
||||
Command output
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
|
||||
# Check if installed
|
||||
if not self.is_installed():
|
||||
error_msg = (
|
||||
f"{self.provider_name} CLI ({self.command_name}) is not installed. "
|
||||
)
|
||||
error_msg += f"Please install it first: https://github.com/anthropics/{self.command_name}"
|
||||
await tool_ctx.error(error_msg)
|
||||
return f"Error: {error_msg}"
|
||||
|
||||
# Check API key if needed
|
||||
if not self.has_api_key():
|
||||
error_msg = f"No API key found for {self.provider_name}. "
|
||||
error_msg += f"Set one of: {', '.join(self.env_vars)}"
|
||||
await tool_ctx.error(error_msg)
|
||||
return f"Error: {error_msg}"
|
||||
|
||||
# Get command arguments
|
||||
cli_args = self.get_cli_args(prompt, **kwargs)
|
||||
|
||||
# Log command
|
||||
await tool_ctx.info(
|
||||
f"Executing {self.provider_name}: {self.command_name} {' '.join(cli_args[:3])}..."
|
||||
)
|
||||
|
||||
try:
|
||||
# Create temp file for prompt if needed
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w", suffix=".txt", delete=False
|
||||
) as f:
|
||||
f.write(prompt)
|
||||
prompt_file = f.name
|
||||
|
||||
# Some CLIs might need the prompt via stdin or file
|
||||
if "--prompt-file" in cli_args:
|
||||
# Replace placeholder with actual file
|
||||
cli_args = [
|
||||
(
|
||||
arg.replace("--prompt-file", prompt_file)
|
||||
if arg == "--prompt-file"
|
||||
else arg
|
||||
)
|
||||
for arg in cli_args
|
||||
]
|
||||
|
||||
# Execute command
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
self.command_name,
|
||||
*cli_args,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
cwd=working_dir or os.getcwd(),
|
||||
)
|
||||
|
||||
# Send prompt via stdin if not using file
|
||||
if "--prompt-file" not in cli_args:
|
||||
stdout, stderr = await asyncio.wait_for(
|
||||
process.communicate(input=prompt.encode()), timeout=timeout
|
||||
)
|
||||
else:
|
||||
stdout, stderr = await asyncio.wait_for(
|
||||
process.communicate(), timeout=timeout
|
||||
)
|
||||
|
||||
# Clean up temp file
|
||||
try:
|
||||
os.unlink(prompt_file)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if process.returncode != 0:
|
||||
error_msg = stderr.decode() if stderr else "Unknown error"
|
||||
await tool_ctx.error(f"{self.provider_name} failed: {error_msg}")
|
||||
return f"Error: {error_msg}"
|
||||
|
||||
result = stdout.decode()
|
||||
await tool_ctx.info(f"{self.provider_name} completed successfully")
|
||||
return result
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
await tool_ctx.error(
|
||||
f"{self.provider_name} timed out after {timeout} seconds"
|
||||
)
|
||||
return f"Error: Command timed out after {timeout} seconds"
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"{self.provider_name} error: {str(e)}")
|
||||
return f"Error: {str(e)}"
|
||||
|
||||
async def call(self, ctx: MCPContext, prompts: str, **kwargs) -> str:
|
||||
"""Execute the CLI agent.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
prompts: The prompt(s) to send
|
||||
**kwargs: Additional arguments
|
||||
|
||||
Returns:
|
||||
Agent response
|
||||
"""
|
||||
return await self.execute_cli(ctx, prompts, **kwargs)
|
||||
@@ -1,543 +0,0 @@
|
||||
"""CLI tool implementations for direct batch execution.
|
||||
|
||||
This module provides CLI tool wrappers that can be used directly in batch operations,
|
||||
including claude (cc), codex, gemini, grok, openhands (oh), hanzo-dev, cline, and aider.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import asyncio
|
||||
from typing import Any, Dict, List, Unpack, Optional, Annotated, TypedDict, final, override
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
# Parameter types for CLI tools
|
||||
Prompt = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="The prompt or command to send to the CLI tool",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
Model = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Optional model override for the CLI tool",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
WorkingDir = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Working directory for the command",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Timeout = Annotated[
|
||||
Optional[int],
|
||||
Field(
|
||||
description="Timeout in seconds for the command",
|
||||
default=300, # 5 minutes default
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class CLIToolParams(TypedDict, total=False):
|
||||
"""Common parameters for CLI tools."""
|
||||
prompt: str
|
||||
model: Optional[str]
|
||||
working_dir: Optional[str]
|
||||
timeout: Optional[int]
|
||||
|
||||
|
||||
class BaseCLITool(BaseTool):
|
||||
"""Base class for CLI tool implementations."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: Optional[PermissionManager] = None,
|
||||
default_model: Optional[str] = None,
|
||||
api_key_env: Optional[str] = None,
|
||||
):
|
||||
"""Initialize CLI tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
default_model: Default model to use
|
||||
api_key_env: Environment variable name for API key
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.default_model = default_model
|
||||
self.api_key_env = api_key_env
|
||||
|
||||
def get_auth_env(self) -> dict[str, str]:
|
||||
"""Get authentication environment variables."""
|
||||
env = os.environ.copy()
|
||||
|
||||
# Add API key if configured
|
||||
if self.api_key_env and self.api_key_env in os.environ:
|
||||
env[self.api_key_env] = os.environ[self.api_key_env]
|
||||
|
||||
# Add Hanzo API key for unified auth
|
||||
if "HANZO_API_KEY" in os.environ:
|
||||
env["HANZO_API_KEY"] = os.environ["HANZO_API_KEY"]
|
||||
|
||||
return env
|
||||
|
||||
async def execute_cli(
|
||||
self,
|
||||
command: list[str],
|
||||
input_text: Optional[str] = None,
|
||||
working_dir: Optional[str] = None,
|
||||
timeout: int = 300,
|
||||
) -> str:
|
||||
"""Execute CLI command with proper error handling.
|
||||
|
||||
Args:
|
||||
command: Command and arguments
|
||||
input_text: Optional stdin input
|
||||
working_dir: Working directory
|
||||
timeout: Timeout in seconds
|
||||
|
||||
Returns:
|
||||
Command output
|
||||
"""
|
||||
try:
|
||||
# Set up environment with auth
|
||||
env = self.get_auth_env()
|
||||
|
||||
# Execute command
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*command,
|
||||
stdin=asyncio.subprocess.PIPE if input_text else None,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
cwd=working_dir,
|
||||
env=env,
|
||||
)
|
||||
|
||||
# Send input and get output
|
||||
stdout, stderr = await asyncio.wait_for(
|
||||
process.communicate(input_text.encode() if input_text else None),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
# Check for errors
|
||||
if process.returncode != 0:
|
||||
error_msg = stderr.decode() if stderr else "Unknown error"
|
||||
return f"Error: {error_msg}"
|
||||
|
||||
return stdout.decode()
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
return f"Error: Command timed out after {timeout} seconds"
|
||||
except Exception as e:
|
||||
return f"Error executing command: {str(e)}"
|
||||
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this tool with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
"""
|
||||
tool_self = self # Create a reference to self for use in the closure
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def tool_wrapper(
|
||||
prompt: str,
|
||||
ctx: Context[Any, Any, Any],
|
||||
model: Optional[str] = None,
|
||||
working_dir: Optional[str] = None,
|
||||
timeout: int = 300,
|
||||
) -> str:
|
||||
result: str = await tool_self.call(
|
||||
ctx, prompt=prompt, model=model, working_dir=working_dir, timeout=timeout
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
class ClaudeCLITool(BaseCLITool):
|
||||
"""Claude CLI tool (also available as 'cc' alias)."""
|
||||
|
||||
def __init__(self, permission_manager: Optional[PermissionManager] = None):
|
||||
super().__init__(
|
||||
permission_manager=permission_manager,
|
||||
default_model="claude-3-5-sonnet-20241022",
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "claude"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "Execute Claude CLI for AI assistance using Anthropic's models"
|
||||
|
||||
async def call(self, ctx: Context[Any, Any, Any], **params: Any) -> str:
|
||||
prompt: str = params.get("prompt", "")
|
||||
model: Optional[str] = params.get("model") or self.default_model
|
||||
working_dir: Optional[str] = params.get("working_dir")
|
||||
timeout: int = params.get("timeout", 300)
|
||||
|
||||
# Build command
|
||||
command: list[str] = ["claude"]
|
||||
if model:
|
||||
command.extend(["--model", model])
|
||||
|
||||
# Execute
|
||||
return await self.execute_cli(
|
||||
command,
|
||||
input_text=prompt,
|
||||
working_dir=working_dir,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
class ClaudeCodeCLITool(ClaudeCLITool):
|
||||
"""Claude Code CLI tool (cc alias)."""
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "cc"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "Claude Code CLI (alias for claude)"
|
||||
|
||||
|
||||
class CodexCLITool(BaseCLITool):
|
||||
"""OpenAI Codex/GPT-4 CLI tool."""
|
||||
|
||||
def __init__(self, permission_manager: Optional[PermissionManager] = None):
|
||||
super().__init__(
|
||||
permission_manager=permission_manager,
|
||||
default_model="gpt-4-turbo",
|
||||
api_key_env="OPENAI_API_KEY",
|
||||
)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "codex"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "Execute OpenAI Codex/GPT-4 CLI for code generation and AI assistance"
|
||||
|
||||
async def call(self, ctx: Context[Any, Any, Any], **params: Any) -> str:
|
||||
prompt: str = params.get("prompt", "")
|
||||
model: Optional[str] = params.get("model") or self.default_model
|
||||
working_dir: Optional[str] = params.get("working_dir")
|
||||
timeout: int = params.get("timeout", 300)
|
||||
|
||||
# Build command (using openai CLI or custom wrapper)
|
||||
command: list[str] = ["openai", "api", "chat.completions.create"]
|
||||
if model:
|
||||
command.extend(["-m", model])
|
||||
command.extend(["-g", "user", prompt])
|
||||
|
||||
# Execute
|
||||
return await self.execute_cli(
|
||||
command,
|
||||
working_dir=working_dir,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
class GeminiCLITool(BaseCLITool):
|
||||
"""Google Gemini CLI tool."""
|
||||
|
||||
def __init__(self, permission_manager: Optional[PermissionManager] = None):
|
||||
super().__init__(
|
||||
permission_manager=permission_manager,
|
||||
default_model="gemini-1.5-pro",
|
||||
api_key_env="GEMINI_API_KEY",
|
||||
)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "gemini"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "Execute Google Gemini CLI for multimodal AI assistance"
|
||||
|
||||
async def call(self, ctx: Context[Any, Any, Any], **params: Any) -> str:
|
||||
prompt: str = params.get("prompt", "")
|
||||
model: Optional[str] = params.get("model") or self.default_model
|
||||
working_dir: Optional[str] = params.get("working_dir")
|
||||
timeout: int = params.get("timeout", 300)
|
||||
|
||||
# Build command
|
||||
command: list[str] = ["gemini"]
|
||||
if model:
|
||||
command.extend(["--model", model])
|
||||
command.append(prompt)
|
||||
|
||||
# Execute
|
||||
return await self.execute_cli(
|
||||
command,
|
||||
working_dir=working_dir,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
class GrokCLITool(BaseCLITool):
|
||||
"""xAI Grok CLI tool."""
|
||||
|
||||
def __init__(self, permission_manager: Optional[PermissionManager] = None):
|
||||
super().__init__(
|
||||
permission_manager=permission_manager,
|
||||
default_model="grok-4",
|
||||
api_key_env="XAI_API_KEY",
|
||||
)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "grok"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "Execute xAI Grok CLI for real-time AI assistance"
|
||||
|
||||
async def call(self, ctx: Context[Any, Any, Any], **params: Any) -> str:
|
||||
prompt: str = params.get("prompt", "")
|
||||
model: Optional[str] = params.get("model") or self.default_model
|
||||
working_dir: Optional[str] = params.get("working_dir")
|
||||
timeout: int = params.get("timeout", 300)
|
||||
|
||||
# Build command
|
||||
command: list[str] = ["grok"]
|
||||
if model:
|
||||
command.extend(["--model", model])
|
||||
command.append(prompt)
|
||||
|
||||
# Execute
|
||||
return await self.execute_cli(
|
||||
command,
|
||||
working_dir=working_dir,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
class OpenHandsCLITool(BaseCLITool):
|
||||
"""OpenHands (OpenDevin) CLI tool."""
|
||||
|
||||
def __init__(self, permission_manager: Optional[PermissionManager] = None):
|
||||
super().__init__(
|
||||
permission_manager=permission_manager,
|
||||
default_model="claude-3-5-sonnet-20241022",
|
||||
api_key_env="OPENAI_API_KEY",
|
||||
)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "openhands"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "Execute OpenHands (OpenDevin) for autonomous coding assistance"
|
||||
|
||||
async def call(self, ctx: Context[Any, Any, Any], **params: Any) -> str:
|
||||
prompt = params.get("prompt", "")
|
||||
model = params.get("model") or self.default_model
|
||||
working_dir: str = params.get("working_dir") or os.getcwd()
|
||||
timeout: int = params.get("timeout", 600) # 10 minutes for OpenHands
|
||||
|
||||
# Build command
|
||||
command: list[str] = ["openhands", "run", prompt]
|
||||
if model:
|
||||
command.extend(["--model", model])
|
||||
command.extend(["--workspace", working_dir])
|
||||
|
||||
# Execute
|
||||
return await self.execute_cli(
|
||||
command,
|
||||
working_dir=working_dir,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
class OpenHandsShortCLITool(OpenHandsCLITool):
|
||||
"""OpenHands CLI tool (oh alias)."""
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "oh"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "OpenHands CLI (alias for openhands)"
|
||||
|
||||
|
||||
class HanzoDevCLITool(BaseCLITool):
|
||||
"""Hanzo Dev AI coding assistant."""
|
||||
|
||||
def __init__(self, permission_manager: Optional[PermissionManager] = None):
|
||||
super().__init__(
|
||||
permission_manager=permission_manager,
|
||||
default_model="claude-3-5-sonnet-20241022",
|
||||
api_key_env="HANZO_API_KEY",
|
||||
)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "hanzo_dev"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "Execute Hanzo Dev for AI-powered code editing and development"
|
||||
|
||||
async def call(self, ctx: Context[Any, Any, Any], **params: Any) -> str:
|
||||
prompt = params.get("prompt", "")
|
||||
model = params.get("model") or self.default_model
|
||||
working_dir: str = params.get("working_dir") or os.getcwd()
|
||||
timeout: int = params.get("timeout", 600)
|
||||
|
||||
# Build command
|
||||
command: list[str] = ["dev"]
|
||||
if model:
|
||||
command.extend(["--model", model])
|
||||
command.extend(["--prompt", prompt])
|
||||
|
||||
# Execute
|
||||
return await self.execute_cli(
|
||||
command,
|
||||
working_dir=working_dir,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
class ClineCLITool(BaseCLITool):
|
||||
"""Cline (formerly Claude Engineer) CLI tool."""
|
||||
|
||||
def __init__(self, permission_manager: Optional[PermissionManager] = None):
|
||||
super().__init__(
|
||||
permission_manager=permission_manager,
|
||||
default_model="claude-3-5-sonnet-20241022",
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "cline"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "Execute Cline for autonomous coding with Claude"
|
||||
|
||||
async def call(self, ctx: Context[Any, Any, Any], **params: Any) -> str:
|
||||
prompt = params.get("prompt", "")
|
||||
working_dir: str = params.get("working_dir") or os.getcwd()
|
||||
timeout: int = params.get("timeout", 600)
|
||||
|
||||
# Build command
|
||||
command: list[str] = ["cline", prompt]
|
||||
command.extend(["--no-interactive"]) # Non-interactive mode for batch
|
||||
|
||||
# Execute
|
||||
return await self.execute_cli(
|
||||
command,
|
||||
working_dir=working_dir,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
class AiderCLITool(BaseCLITool):
|
||||
"""Aider AI pair programming tool."""
|
||||
|
||||
def __init__(self, permission_manager: Optional[PermissionManager] = None):
|
||||
super().__init__(
|
||||
permission_manager=permission_manager,
|
||||
default_model="gpt-4-turbo",
|
||||
api_key_env="OPENAI_API_KEY",
|
||||
)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "aider"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "Execute Aider for AI pair programming"
|
||||
|
||||
async def call(self, ctx: Context[Any, Any, Any], **params: Any) -> str:
|
||||
prompt = params.get("prompt", "")
|
||||
model = params.get("model") or self.default_model
|
||||
working_dir: str = params.get("working_dir") or os.getcwd()
|
||||
timeout: int = params.get("timeout", 600)
|
||||
|
||||
# Build command
|
||||
command: list[str] = ["aider"]
|
||||
if model:
|
||||
command.extend(["--model", model])
|
||||
command.extend(["--message", prompt])
|
||||
command.extend(["--yes"]) # Auto-approve changes
|
||||
command.extend(["--no-stream"]) # No streaming for batch
|
||||
|
||||
# Execute
|
||||
return await self.execute_cli(
|
||||
command,
|
||||
working_dir=working_dir,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
def register_cli_tools(
|
||||
mcp_server: FastMCP,
|
||||
permission_manager: Optional[PermissionManager] = None,
|
||||
) -> list[BaseTool]:
|
||||
"""Register all CLI tools with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
permission_manager: Permission manager for access control
|
||||
|
||||
Returns:
|
||||
List of registered CLI tools
|
||||
"""
|
||||
tools: list[BaseTool] = [
|
||||
ClaudeCLITool(permission_manager),
|
||||
ClaudeCodeCLITool(permission_manager), # cc alias
|
||||
CodexCLITool(permission_manager),
|
||||
GeminiCLITool(permission_manager),
|
||||
GrokCLITool(permission_manager),
|
||||
OpenHandsCLITool(permission_manager),
|
||||
OpenHandsShortCLITool(permission_manager), # oh alias
|
||||
HanzoDevCLITool(permission_manager),
|
||||
ClineCLITool(permission_manager),
|
||||
AiderCLITool(permission_manager),
|
||||
]
|
||||
|
||||
# Register each tool
|
||||
for tool in tools:
|
||||
tool.register(mcp_server)
|
||||
|
||||
return tools
|
||||
|
||||
|
||||
# Export all CLI tool classes
|
||||
__all__ = [
|
||||
"ClaudeCLITool",
|
||||
"ClaudeCodeCLITool",
|
||||
"CodexCLITool",
|
||||
"GeminiCLITool",
|
||||
"GrokCLITool",
|
||||
"OpenHandsCLITool",
|
||||
"OpenHandsShortCLITool",
|
||||
"HanzoDevCLITool",
|
||||
"ClineCLITool",
|
||||
"AiderCLITool",
|
||||
"register_cli_tools",
|
||||
]
|
||||
@@ -1,431 +0,0 @@
|
||||
"""Claude Code and OpenAI Codex authentication management.
|
||||
|
||||
This module provides tools to manage API keys and authentication for
|
||||
Claude Code CLI and OpenAI Codex, allowing separate accounts for swarm agents.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import getpass
|
||||
import subprocess
|
||||
from typing import Any, Dict, List, Tuple, Optional
|
||||
from pathlib import Path
|
||||
from dataclasses import dataclass
|
||||
|
||||
import keyring
|
||||
|
||||
|
||||
@dataclass
|
||||
class APICredential:
|
||||
"""API credential information."""
|
||||
|
||||
provider: str
|
||||
api_key: str
|
||||
model: Optional[str] = None
|
||||
base_url: Optional[str] = None
|
||||
org_id: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
|
||||
|
||||
class CodeAuthManager:
|
||||
"""Manages authentication for Claude Code and other AI coding tools."""
|
||||
|
||||
# Configuration paths
|
||||
CONFIG_DIR = Path.home() / ".hanzo" / "auth"
|
||||
ACCOUNTS_FILE = CONFIG_DIR / "accounts.json"
|
||||
ACTIVE_ACCOUNT_FILE = CONFIG_DIR / "active_account"
|
||||
|
||||
# Environment variable mappings
|
||||
ENV_VARS = {
|
||||
"claude": ["ANTHROPIC_API_KEY", "CLAUDE_API_KEY"],
|
||||
"openai": ["OPENAI_API_KEY"],
|
||||
"azure": ["AZURE_OPENAI_API_KEY", "AZURE_API_KEY"],
|
||||
"deepseek": ["DEEPSEEK_API_KEY"],
|
||||
"google": ["GOOGLE_API_KEY", "GEMINI_API_KEY"],
|
||||
"groq": ["GROQ_API_KEY"],
|
||||
}
|
||||
|
||||
# Default models
|
||||
DEFAULT_MODELS = {
|
||||
"claude": "claude-3-5-sonnet-20241022", # Latest Sonnet
|
||||
"openai": "gpt-4o",
|
||||
"azure": "gpt-4",
|
||||
"deepseek": "deepseek-coder",
|
||||
"google": "gemini-1.5-pro",
|
||||
"groq": "llama3-70b-8192",
|
||||
}
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize auth manager."""
|
||||
self.ensure_config_dir()
|
||||
self._env_backup = {}
|
||||
|
||||
def ensure_config_dir(self):
|
||||
"""Ensure config directory exists."""
|
||||
self.CONFIG_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def get_active_account(self) -> Optional[str]:
|
||||
"""Get the currently active account."""
|
||||
if self.ACTIVE_ACCOUNT_FILE.exists():
|
||||
return self.ACTIVE_ACCOUNT_FILE.read_text().strip()
|
||||
return "default"
|
||||
|
||||
def set_active_account(self, account: str):
|
||||
"""Set the active account."""
|
||||
self.ACTIVE_ACCOUNT_FILE.write_text(account)
|
||||
|
||||
def _load_accounts(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Load all accounts."""
|
||||
if not self.ACCOUNTS_FILE.exists():
|
||||
return {}
|
||||
|
||||
try:
|
||||
with open(self.ACCOUNTS_FILE, "r") as f:
|
||||
return json.load(f)
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
def _save_accounts(self, accounts: Dict[str, Dict[str, Any]]):
|
||||
"""Save accounts."""
|
||||
with open(self.ACCOUNTS_FILE, "w") as f:
|
||||
json.dump(accounts, f, indent=2)
|
||||
|
||||
def list_accounts(self) -> List[str]:
|
||||
"""List all available accounts."""
|
||||
accounts = self._load_accounts()
|
||||
return list(accounts.keys())
|
||||
|
||||
def get_account_info(self, account: str) -> Optional[Dict[str, Any]]:
|
||||
"""Get information about an account."""
|
||||
accounts = self._load_accounts()
|
||||
return accounts.get(account)
|
||||
|
||||
def create_account(
|
||||
self,
|
||||
account: str,
|
||||
provider: str = "claude",
|
||||
api_key: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
) -> Tuple[bool, str]:
|
||||
"""Create a new account.
|
||||
|
||||
Args:
|
||||
account: Account name
|
||||
provider: Provider (claude, openai, etc.)
|
||||
api_key: API key (will prompt if not provided)
|
||||
model: Model to use (defaults to provider default)
|
||||
description: Account description
|
||||
|
||||
Returns:
|
||||
Tuple of (success, message)
|
||||
"""
|
||||
accounts = self._load_accounts()
|
||||
|
||||
if account in accounts:
|
||||
return False, f"Account '{account}' already exists"
|
||||
|
||||
# Get API key if not provided
|
||||
if not api_key:
|
||||
api_key = self._prompt_for_api_key(provider)
|
||||
if not api_key:
|
||||
return False, "No API key provided"
|
||||
|
||||
# Use default model if not specified
|
||||
if not model:
|
||||
model = self.DEFAULT_MODELS.get(provider)
|
||||
|
||||
# Store in keyring for security
|
||||
try:
|
||||
keyring.set_password(f"hanzo-{provider}", account, api_key)
|
||||
except Exception:
|
||||
# Fallback to file storage (less secure)
|
||||
pass
|
||||
|
||||
# Save account info
|
||||
accounts[account] = {
|
||||
"provider": provider,
|
||||
"model": model,
|
||||
"description": description or f"{provider} account",
|
||||
"created_at": os.path.getmtime(__file__),
|
||||
"has_keyring": self._has_keyring_support(),
|
||||
}
|
||||
|
||||
self._save_accounts(accounts)
|
||||
return True, f"Created account '{account}' for {provider}"
|
||||
|
||||
def _prompt_for_api_key(self, provider: str) -> Optional[str]:
|
||||
"""Prompt user for API key."""
|
||||
prompt = f"Enter {provider.upper()} API key: "
|
||||
try:
|
||||
return getpass.getpass(prompt)
|
||||
except KeyboardInterrupt:
|
||||
return None
|
||||
|
||||
def _has_keyring_support(self) -> bool:
|
||||
"""Check if keyring is available."""
|
||||
try:
|
||||
keyring.get_keyring()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def login(self, account: str = "default") -> Tuple[bool, str]:
|
||||
"""Login to an account by setting environment variables.
|
||||
|
||||
Args:
|
||||
account: Account name to login to
|
||||
|
||||
Returns:
|
||||
Tuple of (success, message)
|
||||
"""
|
||||
accounts = self._load_accounts()
|
||||
|
||||
if account not in accounts:
|
||||
return False, f"Account '{account}' not found"
|
||||
|
||||
account_info = accounts[account]
|
||||
provider = account_info["provider"]
|
||||
|
||||
# Get API key from keyring or prompt
|
||||
api_key = None
|
||||
if account_info.get("has_keyring"):
|
||||
try:
|
||||
api_key = keyring.get_password(f"hanzo-{provider}", account)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if not api_key:
|
||||
# Try environment variable
|
||||
for env_var in self.ENV_VARS.get(provider, []):
|
||||
if env_var in os.environ:
|
||||
api_key = os.environ[env_var]
|
||||
break
|
||||
|
||||
if not api_key:
|
||||
api_key = self._prompt_for_api_key(provider)
|
||||
if not api_key:
|
||||
return False, "No API key available"
|
||||
|
||||
# Backup current environment
|
||||
self._backup_environment(provider)
|
||||
|
||||
# Set environment variables
|
||||
for env_var in self.ENV_VARS.get(provider, []):
|
||||
os.environ[env_var] = api_key
|
||||
|
||||
# Set active account
|
||||
self.set_active_account(account)
|
||||
|
||||
# Update shell if using claude command
|
||||
self._update_claude_command(account_info)
|
||||
|
||||
return True, f"Logged in as '{account}' ({provider})"
|
||||
|
||||
def logout(self) -> Tuple[bool, str]:
|
||||
"""Logout by clearing environment variables."""
|
||||
current = self.get_active_account()
|
||||
|
||||
if not current or current == "default":
|
||||
return False, "No active session"
|
||||
|
||||
accounts = self._load_accounts()
|
||||
if current not in accounts:
|
||||
return False, f"Unknown account: {current}"
|
||||
|
||||
provider = accounts[current]["provider"]
|
||||
|
||||
# Clear environment variables
|
||||
for env_var in self.ENV_VARS.get(provider, []):
|
||||
if env_var in os.environ:
|
||||
del os.environ[env_var]
|
||||
|
||||
# Restore backed up environment if any
|
||||
self._restore_environment(provider)
|
||||
|
||||
# Clear active account
|
||||
if self.ACTIVE_ACCOUNT_FILE.exists():
|
||||
self.ACTIVE_ACCOUNT_FILE.unlink()
|
||||
|
||||
return True, f"Logged out from '{current}'"
|
||||
|
||||
def _backup_environment(self, provider: str):
|
||||
"""Backup current environment variables."""
|
||||
for env_var in self.ENV_VARS.get(provider, []):
|
||||
if env_var in os.environ:
|
||||
self._env_backup[env_var] = os.environ[env_var]
|
||||
|
||||
def _restore_environment(self, provider: str):
|
||||
"""Restore backed up environment variables."""
|
||||
for env_var in self.ENV_VARS.get(provider, []):
|
||||
if env_var in self._env_backup:
|
||||
os.environ[env_var] = self._env_backup[env_var]
|
||||
del self._env_backup[env_var]
|
||||
|
||||
def _update_claude_command(self, account_info: Dict[str, Any]):
|
||||
"""Update claude command configuration if needed."""
|
||||
# Check if claude command exists
|
||||
try:
|
||||
result = subprocess.run(["which", "claude"], capture_output=True, text=True)
|
||||
if result.returncode == 0:
|
||||
# Claude command exists, update its config
|
||||
claude_config = Path.home() / ".claude" / "config.json"
|
||||
if claude_config.exists():
|
||||
try:
|
||||
with open(claude_config, "r") as f:
|
||||
config = json.load(f)
|
||||
|
||||
# Update model if specified
|
||||
if account_info.get("model"):
|
||||
config["default_model"] = account_info["model"]
|
||||
|
||||
with open(claude_config, "w") as f:
|
||||
json.dump(config, f, indent=2)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def switch_account(self, account: str) -> Tuple[bool, str]:
|
||||
"""Switch to a different account."""
|
||||
# Logout current
|
||||
self.logout()
|
||||
|
||||
# Login to new account
|
||||
return self.login(account)
|
||||
|
||||
def create_agent_account(
|
||||
self,
|
||||
agent_id: str,
|
||||
provider: str = "claude",
|
||||
parent_account: Optional[str] = None,
|
||||
) -> Tuple[bool, str]:
|
||||
"""Create an account for a swarm agent.
|
||||
|
||||
Args:
|
||||
agent_id: Unique agent identifier
|
||||
provider: AI provider
|
||||
parent_account: Parent account to clone from
|
||||
|
||||
Returns:
|
||||
Tuple of (success, account_name)
|
||||
"""
|
||||
agent_account = f"agent_{agent_id}"
|
||||
|
||||
# If parent account specified, clone its credentials
|
||||
if parent_account:
|
||||
parent_info = self.get_account_info(parent_account)
|
||||
if not parent_info:
|
||||
return False, f"Parent account '{parent_account}' not found"
|
||||
|
||||
# Get parent API key
|
||||
api_key = None
|
||||
if parent_info.get("has_keyring"):
|
||||
try:
|
||||
api_key = keyring.get_password(
|
||||
f"hanzo-{parent_info['provider']}", parent_account
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if api_key:
|
||||
success, msg = self.create_account(
|
||||
agent_account,
|
||||
provider=parent_info["provider"],
|
||||
api_key=api_key,
|
||||
model=parent_info.get("model"),
|
||||
description=f"Agent account (parent: {parent_account})",
|
||||
)
|
||||
if success:
|
||||
return True, agent_account
|
||||
|
||||
# Create with current environment
|
||||
for env_var in self.ENV_VARS.get(provider, []):
|
||||
if env_var in os.environ:
|
||||
success, msg = self.create_account(
|
||||
agent_account,
|
||||
provider=provider,
|
||||
api_key=os.environ[env_var],
|
||||
model=self.DEFAULT_MODELS.get(provider),
|
||||
description=f"Agent account for {agent_id}",
|
||||
)
|
||||
if success:
|
||||
return True, agent_account
|
||||
|
||||
return False, "No credentials available for agent"
|
||||
|
||||
def get_agent_credentials(self, agent_id: str) -> Optional[APICredential]:
|
||||
"""Get credentials for an agent.
|
||||
|
||||
Args:
|
||||
agent_id: Agent identifier
|
||||
|
||||
Returns:
|
||||
APICredential if found
|
||||
"""
|
||||
agent_account = f"agent_{agent_id}"
|
||||
account_info = self.get_account_info(agent_account)
|
||||
|
||||
if not account_info:
|
||||
return None
|
||||
|
||||
# Get API key
|
||||
api_key = None
|
||||
provider = account_info["provider"]
|
||||
|
||||
if account_info.get("has_keyring"):
|
||||
try:
|
||||
api_key = keyring.get_password(f"hanzo-{provider}", agent_account)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if not api_key:
|
||||
# Try current environment
|
||||
for env_var in self.ENV_VARS.get(provider, []):
|
||||
if env_var in os.environ:
|
||||
api_key = os.environ[env_var]
|
||||
break
|
||||
|
||||
if not api_key:
|
||||
return None
|
||||
|
||||
return APICredential(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
model=account_info.get("model"),
|
||||
description=account_info.get("description"),
|
||||
)
|
||||
|
||||
|
||||
# Update swarm tool to use latest Sonnet
|
||||
def get_latest_claude_model() -> str:
|
||||
"""Get the latest Claude model identifier."""
|
||||
# As of the knowledge cutoff, this is the latest Sonnet
|
||||
# In production, this could query an API for the latest model
|
||||
return "claude-3-5-sonnet-20241022"
|
||||
|
||||
|
||||
# Token counting using tiktoken (same as current implementation)
|
||||
def count_tokens_streaming(text_stream) -> int:
|
||||
"""Count tokens in a streaming fashion.
|
||||
|
||||
This uses the same tiktoken approach as the truncate module,
|
||||
but processes text as it streams.
|
||||
"""
|
||||
import tiktoken
|
||||
|
||||
try:
|
||||
# Use cl100k_base encoding (Claude/GPT-4 compatible)
|
||||
encoding = tiktoken.get_encoding("cl100k_base")
|
||||
except Exception:
|
||||
# Fallback to simple estimation
|
||||
return len(text_stream) // 4
|
||||
|
||||
total_tokens = 0
|
||||
for chunk in text_stream:
|
||||
if isinstance(chunk, str):
|
||||
total_tokens += len(encoding.encode(chunk))
|
||||
elif isinstance(chunk, bytes):
|
||||
total_tokens += len(encoding.encode(chunk.decode("utf-8", errors="ignore")))
|
||||
|
||||
return total_tokens
|
||||
@@ -1,195 +0,0 @@
|
||||
"""Claude Code authentication tool.
|
||||
|
||||
This tool manages API keys and accounts for Claude Code and other AI coding tools.
|
||||
"""
|
||||
|
||||
from typing import Unpack, Optional, TypedDict, final, override
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.agent.code_auth import CodeAuthManager
|
||||
|
||||
|
||||
class CodeAuthParams(TypedDict, total=False):
|
||||
"""Parameters for code auth tool."""
|
||||
|
||||
action: str
|
||||
account: Optional[str]
|
||||
provider: Optional[str]
|
||||
api_key: Optional[str]
|
||||
model: Optional[str]
|
||||
description: Optional[str]
|
||||
agent_id: Optional[str]
|
||||
parent_account: Optional[str]
|
||||
|
||||
|
||||
@final
|
||||
class CodeAuthTool(BaseTool):
|
||||
"""Tool for managing Claude Code authentication and API keys."""
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "code_auth"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Manage Claude Code and AI provider authentication.
|
||||
|
||||
Actions:
|
||||
- status: Show current login status
|
||||
- list: List all accounts
|
||||
- create: Create a new account
|
||||
- login: Login to an account
|
||||
- logout: Logout current account
|
||||
- switch: Switch between accounts
|
||||
- agent: Create/get agent account
|
||||
|
||||
Examples:
|
||||
code_auth status
|
||||
code_auth list
|
||||
code_auth create --account work --provider claude
|
||||
code_auth login --account work
|
||||
code_auth logout
|
||||
code_auth switch --account personal
|
||||
code_auth agent --agent_id swarm_1 --parent_account work
|
||||
|
||||
Providers: claude, openai, azure, deepseek, google, groq"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the code auth tool."""
|
||||
self.auth_manager = CodeAuthManager()
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[CodeAuthParams],
|
||||
) -> str:
|
||||
"""Execute the code auth tool.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Result message
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
action = params.get("action", "status")
|
||||
|
||||
if action == "status":
|
||||
current = self.auth_manager.get_active_account()
|
||||
if current:
|
||||
info = self.auth_manager.get_account_info(current)
|
||||
if info:
|
||||
return f"Logged in as: {current} ({info['provider']})"
|
||||
return "Not logged in"
|
||||
|
||||
elif action == "list":
|
||||
accounts = self.auth_manager.list_accounts()
|
||||
if not accounts:
|
||||
return "No accounts configured"
|
||||
|
||||
current = self.auth_manager.get_active_account()
|
||||
lines = ["Configured accounts:"]
|
||||
for account in accounts:
|
||||
info = self.auth_manager.get_account_info(account)
|
||||
marker = " (active)" if account == current else ""
|
||||
lines.append(f" - {account}: {info['provider']}{marker}")
|
||||
return "\n".join(lines)
|
||||
|
||||
elif action == "create":
|
||||
account = params.get("account")
|
||||
if not account:
|
||||
return "Error: account name required"
|
||||
|
||||
provider = params.get("provider", "claude")
|
||||
api_key = params.get("api_key")
|
||||
model = params.get("model")
|
||||
description = params.get("description")
|
||||
|
||||
success, msg = self.auth_manager.create_account(
|
||||
account, provider, api_key, model, description
|
||||
)
|
||||
return msg
|
||||
|
||||
elif action == "login":
|
||||
account = params.get("account", "default")
|
||||
success, msg = self.auth_manager.login(account)
|
||||
return msg
|
||||
|
||||
elif action == "logout":
|
||||
success, msg = self.auth_manager.logout()
|
||||
return msg
|
||||
|
||||
elif action == "switch":
|
||||
account = params.get("account")
|
||||
if not account:
|
||||
return "Error: account name required"
|
||||
|
||||
success, msg = self.auth_manager.switch_account(account)
|
||||
return msg
|
||||
|
||||
elif action == "agent":
|
||||
agent_id = params.get("agent_id")
|
||||
if not agent_id:
|
||||
return "Error: agent_id required"
|
||||
|
||||
provider = params.get("provider", "claude")
|
||||
parent_account = params.get("parent_account")
|
||||
|
||||
# Try to create agent account
|
||||
success, result = self.auth_manager.create_agent_account(
|
||||
agent_id, provider, parent_account
|
||||
)
|
||||
|
||||
if success:
|
||||
# Get credentials
|
||||
creds = self.auth_manager.get_agent_credentials(agent_id)
|
||||
if creds:
|
||||
return f"Agent account ready: {result} ({creds.provider})"
|
||||
else:
|
||||
return f"Agent account created but no credentials: {result}"
|
||||
else:
|
||||
return f"Failed to create agent account: {result}"
|
||||
|
||||
else:
|
||||
return f"Unknown action: {action}. Use: status, list, create, login, logout, switch, agent"
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def code_auth(
|
||||
ctx: MCPContext,
|
||||
action: str = "status",
|
||||
account: Optional[str] = None,
|
||||
provider: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
agent_id: Optional[str] = None,
|
||||
parent_account: Optional[str] = None,
|
||||
) -> str:
|
||||
return await tool_self.call(
|
||||
ctx,
|
||||
action=action,
|
||||
account=account,
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
description=description,
|
||||
agent_id=agent_id,
|
||||
parent_account=parent_account,
|
||||
)
|
||||
@@ -1,124 +0,0 @@
|
||||
"""OpenAI Codex CLI agent tool.
|
||||
|
||||
This tool provides integration with OpenAI's CLI (openai command),
|
||||
allowing programmatic execution of GPT-4 and other models for code tasks.
|
||||
"""
|
||||
|
||||
from typing import List, Optional, final, override
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.agent.cli_agent_base import CLIAgentBase
|
||||
|
||||
|
||||
@final
|
||||
class CodexCLITool(CLIAgentBase):
|
||||
"""Tool for executing OpenAI CLI (formerly Codex)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
model: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize Codex CLI tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
model: Optional model override (defaults to gpt-4o)
|
||||
**kwargs: Additional arguments
|
||||
"""
|
||||
super().__init__(
|
||||
permission_manager=permission_manager,
|
||||
command_name="openai",
|
||||
provider_name="OpenAI",
|
||||
default_model=model or "gpt-4o",
|
||||
env_vars=["OPENAI_API_KEY"],
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "codex_cli"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Execute OpenAI CLI for code generation and analysis.
|
||||
|
||||
This tool runs the OpenAI CLI (openai command) for code generation,
|
||||
completion, and analysis tasks. It uses GPT-4o by default but supports
|
||||
all OpenAI models.
|
||||
|
||||
Features:
|
||||
- GPT-4 and GPT-4o for advanced reasoning
|
||||
- Code generation and completion
|
||||
- Multi-modal support (with gpt-4-vision)
|
||||
- Function calling capabilities
|
||||
|
||||
Usage:
|
||||
codex_cli(prompts="Generate a Python function to sort a binary tree")
|
||||
codex_cli(prompts="Explain this code and suggest improvements", model="gpt-4-turbo")
|
||||
|
||||
Requirements:
|
||||
- OpenAI CLI must be installed: pip install openai
|
||||
- OPENAI_API_KEY environment variable
|
||||
"""
|
||||
|
||||
@override
|
||||
def get_cli_args(self, prompt: str, **kwargs) -> List[str]:
|
||||
"""Get CLI arguments for OpenAI.
|
||||
|
||||
Args:
|
||||
prompt: The prompt to send
|
||||
**kwargs: Additional arguments (model, temperature, etc.)
|
||||
|
||||
Returns:
|
||||
List of command arguments
|
||||
"""
|
||||
args = ["api", "chat.completions.create"]
|
||||
|
||||
# Add model
|
||||
model = kwargs.get("model", self.default_model)
|
||||
args.extend(["-m", model])
|
||||
|
||||
# Add temperature if specified
|
||||
if "temperature" in kwargs:
|
||||
args.extend(["--temperature", str(kwargs["temperature"])])
|
||||
|
||||
# Add max tokens if specified
|
||||
if "max_tokens" in kwargs:
|
||||
args.extend(["--max-tokens", str(kwargs["max_tokens"])])
|
||||
|
||||
# Add the prompt as a message
|
||||
args.extend(["-g", prompt])
|
||||
|
||||
return args
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def codex_cli(
|
||||
ctx: MCPContext,
|
||||
prompts: str,
|
||||
model: Optional[str] = None,
|
||||
temperature: Optional[float] = None,
|
||||
max_tokens: Optional[int] = None,
|
||||
working_dir: Optional[str] = None,
|
||||
) -> str:
|
||||
return await tool_self.call(
|
||||
ctx,
|
||||
prompts=prompts,
|
||||
model=model,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
working_dir=working_dir,
|
||||
)
|
||||
@@ -1,389 +0,0 @@
|
||||
"""Critic tool for agents to request critical review from main loop."""
|
||||
|
||||
from enum import Enum
|
||||
from typing import List, Optional, override
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
|
||||
|
||||
class ReviewType(Enum):
|
||||
"""Types of review requests."""
|
||||
|
||||
CODE_QUALITY = "code_quality"
|
||||
CORRECTNESS = "correctness"
|
||||
PERFORMANCE = "performance"
|
||||
SECURITY = "security"
|
||||
COMPLETENESS = "completeness"
|
||||
BEST_PRACTICES = "best_practices"
|
||||
GENERAL = "general"
|
||||
|
||||
|
||||
class CriticTool(BaseTool):
|
||||
"""Tool for agents to request critical review from the main loop."""
|
||||
|
||||
name = "critic"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Request critical review and feedback from the main loop.
|
||||
|
||||
Use this tool to get automated critical analysis of your work. The main loop will:
|
||||
- Review your implementation for bugs, edge cases, and improvements
|
||||
- Check for security issues and best practices
|
||||
- Suggest performance optimizations
|
||||
- Ensure completeness and correctness
|
||||
- Provide actionable feedback for improvements
|
||||
|
||||
Parameters:
|
||||
- review_type: Type of review (CODE_QUALITY, CORRECTNESS, PERFORMANCE, SECURITY, COMPLETENESS, BEST_PRACTICES, GENERAL)
|
||||
- work_description: Clear description of what you've done
|
||||
- code_snippets: Optional code snippets to review (as a list of strings)
|
||||
- file_paths: Optional list of file paths you've modified
|
||||
- specific_concerns: Optional specific areas you want reviewed
|
||||
|
||||
The critic will provide harsh but constructive feedback to ensure high quality.
|
||||
|
||||
Example:
|
||||
critic(
|
||||
review_type="CODE_QUALITY",
|
||||
work_description="Added import statements to fix undefined symbols in Go files",
|
||||
code_snippets=["import (\n \"fmt\"\n \"github.com/luxfi/node/common\"\n)"],
|
||||
file_paths=["/path/to/atomic.go", "/path/to/network.go"],
|
||||
specific_concerns="Are the imports in the correct format and location?"
|
||||
)"""
|
||||
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
review_type: str,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]] = None,
|
||||
file_paths: Optional[List[str]] = None,
|
||||
specific_concerns: Optional[str] = None,
|
||||
) -> str:
|
||||
"""This is a placeholder - actual implementation happens in AgentTool."""
|
||||
# This tool is handled specially in the agent execution
|
||||
return f"Critic review requested for: {work_description}"
|
||||
|
||||
def register(self, server: FastMCP) -> None:
|
||||
"""Register the tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@server.tool(name=self.name, description=self.description)
|
||||
async def critic(
|
||||
ctx: MCPContext,
|
||||
review_type: str,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]] = None,
|
||||
file_paths: Optional[List[str]] = None,
|
||||
specific_concerns: Optional[str] = None,
|
||||
) -> str:
|
||||
return await tool_self.call(
|
||||
ctx,
|
||||
review_type,
|
||||
work_description,
|
||||
code_snippets,
|
||||
file_paths,
|
||||
specific_concerns,
|
||||
)
|
||||
|
||||
|
||||
class AutoCritic:
|
||||
"""Automated critic that provides harsh but constructive feedback."""
|
||||
|
||||
def __init__(self):
|
||||
self.review_patterns = {
|
||||
ReviewType.CODE_QUALITY: self._review_code_quality,
|
||||
ReviewType.CORRECTNESS: self._review_correctness,
|
||||
ReviewType.PERFORMANCE: self._review_performance,
|
||||
ReviewType.SECURITY: self._review_security,
|
||||
ReviewType.COMPLETENESS: self._review_completeness,
|
||||
ReviewType.BEST_PRACTICES: self._review_best_practices,
|
||||
ReviewType.GENERAL: self._review_general,
|
||||
}
|
||||
|
||||
def review(
|
||||
self,
|
||||
review_type: ReviewType,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]] = None,
|
||||
file_paths: Optional[List[str]] = None,
|
||||
specific_concerns: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Perform automated critical review."""
|
||||
review_func = self.review_patterns.get(review_type, self._review_general)
|
||||
return review_func(
|
||||
work_description, code_snippets, file_paths, specific_concerns
|
||||
)
|
||||
|
||||
def _review_code_quality(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
specific_concerns: Optional[str],
|
||||
) -> str:
|
||||
"""Review code quality aspects."""
|
||||
issues = []
|
||||
suggestions = []
|
||||
|
||||
# Check for common code quality issues
|
||||
if code_snippets:
|
||||
for snippet in code_snippets:
|
||||
# Check for proper error handling
|
||||
if "error" in snippet.lower() and "if err" not in snippet:
|
||||
issues.append(
|
||||
"❌ Missing error handling - always check errors in Go"
|
||||
)
|
||||
|
||||
# Check for magic numbers
|
||||
if any(char.isdigit() for char in snippet) and "const" not in snippet:
|
||||
suggestions.append(
|
||||
"💡 Consider extracting magic numbers to named constants"
|
||||
)
|
||||
|
||||
# Check for proper imports
|
||||
if "import" in snippet:
|
||||
if '"fmt"' in snippet and snippet.count("fmt.") == 0:
|
||||
issues.append("❌ Unused import 'fmt' - remove unused imports")
|
||||
if not snippet.strip().endswith(")") and "import (" in snippet:
|
||||
issues.append("❌ Import block not properly closed")
|
||||
|
||||
# General quality checks
|
||||
if "fix" in work_description.lower():
|
||||
suggestions.append("💡 Ensure you've tested the fix thoroughly")
|
||||
suggestions.append("💡 Consider edge cases and error scenarios")
|
||||
|
||||
if file_paths and len(file_paths) > 5:
|
||||
suggestions.append(
|
||||
"💡 Large number of files modified - consider breaking into smaller PRs"
|
||||
)
|
||||
|
||||
# Build response
|
||||
response = "🔍 CODE QUALITY REVIEW:\n\n"
|
||||
|
||||
if issues:
|
||||
response += "Issues Found:\n" + "\n".join(issues) + "\n\n"
|
||||
else:
|
||||
response += "✅ No major code quality issues detected.\n\n"
|
||||
|
||||
if suggestions:
|
||||
response += (
|
||||
"Suggestions for Improvement:\n" + "\n".join(suggestions) + "\n\n"
|
||||
)
|
||||
|
||||
if specific_concerns:
|
||||
response += f"Regarding your concern: '{specific_concerns}'\n"
|
||||
if "import" in specific_concerns.lower():
|
||||
response += "→ Imports look properly formatted. Ensure they're in the standard order: stdlib, external, internal.\n"
|
||||
|
||||
response += "\nOverall: " + (
|
||||
"⚠️ Address the issues before proceeding."
|
||||
if issues
|
||||
else "✅ Good work, but always room for improvement!"
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
def _review_correctness(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
specific_concerns: Optional[str],
|
||||
) -> str:
|
||||
"""Review correctness aspects."""
|
||||
return """🔍 CORRECTNESS REVIEW:
|
||||
|
||||
Critical Questions:
|
||||
❓ Have you verified the changes compile without errors?
|
||||
❓ Do the changes actually fix the reported issue?
|
||||
❓ Have you introduced any new bugs or regressions?
|
||||
❓ Are all edge cases handled properly?
|
||||
|
||||
Specific Checks:
|
||||
- If fixing imports: Verify the import paths are correct for the project
|
||||
- If modifying logic: Ensure the logic is sound and handles all cases
|
||||
- If refactoring: Confirm behavior is preserved
|
||||
|
||||
⚠️ Remember: Working code > elegant code. Make sure it works first!"""
|
||||
|
||||
def _review_performance(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
specific_concerns: Optional[str],
|
||||
) -> str:
|
||||
"""Review performance aspects."""
|
||||
return """🔍 PERFORMANCE REVIEW:
|
||||
|
||||
Performance Considerations:
|
||||
- Are you doing any operations in loops that could be moved outside?
|
||||
- Are there any unnecessary allocations or copies?
|
||||
- Could any synchronous operations be made concurrent?
|
||||
- Are you caching results that might be reused?
|
||||
|
||||
For file operations:
|
||||
- Consider batch operations over individual ones
|
||||
- Use buffered I/O for large files
|
||||
- Avoid reading entire files into memory if possible
|
||||
|
||||
💡 Remember: Premature optimization is evil, but obvious inefficiencies should be fixed."""
|
||||
|
||||
def _review_security(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
specific_concerns: Optional[str],
|
||||
) -> str:
|
||||
"""Review security aspects."""
|
||||
return """🔍 SECURITY REVIEW:
|
||||
|
||||
Security Checklist:
|
||||
🔐 No hardcoded secrets or credentials
|
||||
🔐 All user inputs are validated/sanitized
|
||||
🔐 File paths are properly validated
|
||||
🔐 No SQL injection vulnerabilities
|
||||
🔐 Proper access control checks
|
||||
🔐 Sensitive data is not logged
|
||||
|
||||
⚠️ If in doubt, err on the side of caution!"""
|
||||
|
||||
def _review_completeness(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
specific_concerns: Optional[str],
|
||||
) -> str:
|
||||
"""Review completeness aspects."""
|
||||
tasks_mentioned = work_description.lower()
|
||||
|
||||
response = "🔍 COMPLETENESS REVIEW:\n\n"
|
||||
|
||||
if "fix" in tasks_mentioned and "test" not in tasks_mentioned:
|
||||
response += (
|
||||
"❌ No mention of tests - have you verified the fix with tests?\n"
|
||||
)
|
||||
|
||||
if "import" in tasks_mentioned:
|
||||
response += "✓ Import fixes mentioned\n"
|
||||
response += "❓ Have you checked for other files with similar issues?\n"
|
||||
response += "❓ Are all undefined symbols now resolved?\n"
|
||||
|
||||
if file_paths:
|
||||
response += f"\n✓ Modified {len(file_paths)} files\n"
|
||||
response += "❓ Are there any related files that also need updates?\n"
|
||||
|
||||
response += "\n💡 Completeness means not just fixing the immediate issue, but considering the broader impact."
|
||||
|
||||
return response
|
||||
|
||||
def _review_best_practices(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
specific_concerns: Optional[str],
|
||||
) -> str:
|
||||
"""Review best practices."""
|
||||
return """🔍 BEST PRACTICES REVIEW:
|
||||
|
||||
Go Best Practices (if applicable):
|
||||
✓ Imports are grouped: stdlib, external, internal
|
||||
✓ Error handling follows Go idioms
|
||||
✓ Variable names are clear and idiomatic
|
||||
✓ Comments explain why, not what
|
||||
✓ Functions do one thing well
|
||||
|
||||
General Best Practices:
|
||||
✓ Code is self-documenting
|
||||
✓ DRY principle is followed
|
||||
✓ SOLID principles are respected
|
||||
✓ Changes are minimal and focused
|
||||
|
||||
💡 Good code is code that others (including future you) can understand and modify."""
|
||||
|
||||
def _review_general(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
specific_concerns: Optional[str],
|
||||
) -> str:
|
||||
"""General review covering multiple aspects."""
|
||||
response = "🔍 GENERAL CRITICAL REVIEW:\n\n"
|
||||
response += f"Work Description: {work_description}\n\n"
|
||||
|
||||
# Quick assessment
|
||||
response += "Quick Assessment:\n"
|
||||
|
||||
if "fix" in work_description.lower():
|
||||
response += "- Type: Bug fix / Error resolution\n"
|
||||
response += "- Critical: Ensure the fix is complete and tested\n"
|
||||
elif "add" in work_description.lower():
|
||||
response += "- Type: Feature addition\n"
|
||||
response += "- Critical: Ensure no regressions introduced\n"
|
||||
elif "refactor" in work_description.lower():
|
||||
response += "- Type: Code refactoring\n"
|
||||
response += "- Critical: Ensure behavior is preserved\n"
|
||||
|
||||
if file_paths:
|
||||
response += f"- Scope: {len(file_paths)} files affected\n"
|
||||
if len(file_paths) > 10:
|
||||
response += "- ⚠️ Large scope - consider breaking down\n"
|
||||
|
||||
response += "\nCritical Questions:\n"
|
||||
response += "1. Is this the minimal change needed?\n"
|
||||
response += "2. Have you considered all edge cases?\n"
|
||||
response += "3. Will this work in production?\n"
|
||||
response += "4. Is there a simpler solution?\n"
|
||||
|
||||
if specific_concerns:
|
||||
response += f"\nYour Concern: {specific_concerns}\n"
|
||||
response += "→ Valid concern. Double-check this area carefully.\n"
|
||||
|
||||
response += "\n🎯 Bottom Line: Good work needs critical thinking. Question everything, verify everything."
|
||||
|
||||
return response
|
||||
|
||||
|
||||
class CriticProtocol:
|
||||
"""Protocol for critic interactions."""
|
||||
|
||||
def __init__(self):
|
||||
self.auto_critic = AutoCritic()
|
||||
self.review_count = 0
|
||||
self.max_reviews = 2 # Allow up to 2 reviews per task
|
||||
|
||||
def request_review(
|
||||
self,
|
||||
review_type: str,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]] = None,
|
||||
file_paths: Optional[List[str]] = None,
|
||||
specific_concerns: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Request a critical review."""
|
||||
if self.review_count >= self.max_reviews:
|
||||
return "❌ Review limit exceeded. Time to move forward with what you have."
|
||||
|
||||
self.review_count += 1
|
||||
|
||||
try:
|
||||
review_enum = ReviewType[review_type.upper()]
|
||||
except KeyError:
|
||||
review_enum = ReviewType.GENERAL
|
||||
|
||||
review = self.auto_critic.review(
|
||||
review_enum, work_description, code_snippets, file_paths, specific_concerns
|
||||
)
|
||||
|
||||
return f"Review {self.review_count}/{self.max_reviews}:\n\n{review}"
|
||||
@@ -1,129 +0,0 @@
|
||||
"""Google Gemini CLI agent tool.
|
||||
|
||||
This tool provides integration with Google's Gemini CLI,
|
||||
allowing programmatic execution of Gemini models for code tasks.
|
||||
"""
|
||||
|
||||
from typing import List, Optional, final, override
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.agent.cli_agent_base import CLIAgentBase
|
||||
|
||||
|
||||
@final
|
||||
class GeminiCLITool(CLIAgentBase):
|
||||
"""Tool for executing Google Gemini CLI."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
model: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize Gemini CLI tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
model: Optional model override (defaults to gemini-1.5-pro)
|
||||
**kwargs: Additional arguments
|
||||
"""
|
||||
super().__init__(
|
||||
permission_manager=permission_manager,
|
||||
command_name="gemini",
|
||||
provider_name="Google Gemini",
|
||||
default_model=model or "gemini-1.5-pro",
|
||||
env_vars=["GOOGLE_API_KEY", "GEMINI_API_KEY"],
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "gemini_cli"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Execute Google Gemini CLI for code tasks.
|
||||
|
||||
This tool runs the Google Gemini CLI for code generation, analysis,
|
||||
and multi-modal tasks. It uses Gemini 1.5 Pro by default.
|
||||
|
||||
Features:
|
||||
- Gemini 1.5 Pro with 2M token context window
|
||||
- Gemini 1.5 Flash for faster responses
|
||||
- Multi-modal capabilities (code + images)
|
||||
- Advanced reasoning and analysis
|
||||
|
||||
Usage:
|
||||
gemini_cli(prompts="Create a React component for a data table")
|
||||
gemini_cli(prompts="Analyze this code for security vulnerabilities", model="gemini-1.5-flash")
|
||||
|
||||
Requirements:
|
||||
- Gemini CLI must be installed
|
||||
- GOOGLE_API_KEY or GEMINI_API_KEY environment variable
|
||||
"""
|
||||
|
||||
@override
|
||||
def get_cli_args(self, prompt: str, **kwargs) -> List[str]:
|
||||
"""Get CLI arguments for Gemini.
|
||||
|
||||
Args:
|
||||
prompt: The prompt to send
|
||||
**kwargs: Additional arguments (model, temperature, etc.)
|
||||
|
||||
Returns:
|
||||
List of command arguments
|
||||
"""
|
||||
args = ["generate"]
|
||||
|
||||
# Add model
|
||||
model = kwargs.get("model", self.default_model)
|
||||
args.extend(["--model", model])
|
||||
|
||||
# Add temperature if specified
|
||||
if "temperature" in kwargs:
|
||||
args.extend(["--temperature", str(kwargs["temperature"])])
|
||||
|
||||
# Add max tokens if specified
|
||||
if "max_tokens" in kwargs:
|
||||
args.extend(["--max-output-tokens", str(kwargs["max_tokens"])])
|
||||
|
||||
# Add safety settings if needed
|
||||
if kwargs.get("safety_settings"):
|
||||
args.extend(["--safety-settings", kwargs["safety_settings"]])
|
||||
|
||||
# Add the prompt
|
||||
args.extend(["--prompt", prompt])
|
||||
|
||||
return args
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def gemini_cli(
|
||||
ctx: MCPContext,
|
||||
prompts: str,
|
||||
model: Optional[str] = None,
|
||||
temperature: Optional[float] = None,
|
||||
max_tokens: Optional[int] = None,
|
||||
working_dir: Optional[str] = None,
|
||||
safety_settings: Optional[str] = None,
|
||||
) -> str:
|
||||
return await tool_self.call(
|
||||
ctx,
|
||||
prompts=prompts,
|
||||
model=model,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
working_dir=working_dir,
|
||||
safety_settings=safety_settings,
|
||||
)
|
||||
@@ -1,129 +0,0 @@
|
||||
"""Grok CLI agent tool.
|
||||
|
||||
This tool provides integration with xAI's Grok CLI,
|
||||
allowing programmatic execution of Grok models for code tasks.
|
||||
"""
|
||||
|
||||
from typing import List, Optional, final, override
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.agent.cli_agent_base import CLIAgentBase
|
||||
|
||||
|
||||
@final
|
||||
class GrokCLITool(CLIAgentBase):
|
||||
"""Tool for executing Grok CLI."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
model: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize Grok CLI tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
model: Optional model override (defaults to grok-2)
|
||||
**kwargs: Additional arguments
|
||||
"""
|
||||
super().__init__(
|
||||
permission_manager=permission_manager,
|
||||
command_name="grok",
|
||||
provider_name="xAI Grok",
|
||||
default_model=model or "grok-2",
|
||||
env_vars=["XAI_API_KEY", "GROK_API_KEY"],
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "grok_cli"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Execute xAI Grok CLI for code tasks.
|
||||
|
||||
This tool runs the Grok CLI for code generation, analysis,
|
||||
and reasoning tasks. It uses Grok-2 by default.
|
||||
|
||||
Features:
|
||||
- Grok-2 with advanced reasoning capabilities
|
||||
- Real-time information access
|
||||
- Code generation and analysis
|
||||
- Humor and personality in responses
|
||||
|
||||
Usage:
|
||||
grok_cli(prompts="Write a Python web scraper with async support")
|
||||
grok_cli(prompts="Explain quantum computing like I'm a programmer", model="grok-1")
|
||||
|
||||
Requirements:
|
||||
- Grok CLI must be installed
|
||||
- XAI_API_KEY or GROK_API_KEY environment variable
|
||||
"""
|
||||
|
||||
@override
|
||||
def get_cli_args(self, prompt: str, **kwargs) -> List[str]:
|
||||
"""Get CLI arguments for Grok.
|
||||
|
||||
Args:
|
||||
prompt: The prompt to send
|
||||
**kwargs: Additional arguments (model, temperature, etc.)
|
||||
|
||||
Returns:
|
||||
List of command arguments
|
||||
"""
|
||||
args = ["chat"]
|
||||
|
||||
# Add model
|
||||
model = kwargs.get("model", self.default_model)
|
||||
args.extend(["--model", model])
|
||||
|
||||
# Add temperature if specified
|
||||
if "temperature" in kwargs:
|
||||
args.extend(["--temperature", str(kwargs["temperature"])])
|
||||
|
||||
# Add max tokens if specified
|
||||
if "max_tokens" in kwargs:
|
||||
args.extend(["--max-tokens", str(kwargs["max_tokens"])])
|
||||
|
||||
# Add system prompt if specified
|
||||
if "system_prompt" in kwargs:
|
||||
args.extend(["--system", kwargs["system_prompt"]])
|
||||
|
||||
# Add the prompt
|
||||
args.append(prompt)
|
||||
|
||||
return args
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def grok_cli(
|
||||
ctx: MCPContext,
|
||||
prompts: str,
|
||||
model: Optional[str] = None,
|
||||
temperature: Optional[float] = None,
|
||||
max_tokens: Optional[int] = None,
|
||||
working_dir: Optional[str] = None,
|
||||
system_prompt: Optional[str] = None,
|
||||
) -> str:
|
||||
return await tool_self.call(
|
||||
ctx,
|
||||
prompts=prompts,
|
||||
model=model,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
working_dir=working_dir,
|
||||
system_prompt=system_prompt,
|
||||
)
|
||||
@@ -1,645 +0,0 @@
|
||||
"""I Ching tool for creative problem solving using Hanzo principles."""
|
||||
|
||||
import random
|
||||
from enum import Enum
|
||||
from typing import List, override
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
|
||||
|
||||
class HanzoPrinciple(Enum):
|
||||
"""Hanzo principles organized by category."""
|
||||
|
||||
# Empathy
|
||||
AUTONOMY = ("Autonomy", "Trust fully; freedom fuels genius", "🦅")
|
||||
BALANCE = ("Balance", "Steady wins; burnout loses every time", "⚖️")
|
||||
CUSTOMER_OBSESSION = (
|
||||
"Customer Obsession",
|
||||
"Coach relentlessly; their victories yours",
|
||||
"🎓",
|
||||
)
|
||||
HUMILITY = ("Humility", "Quiet confidence; greatness emerges naturally", "🧘")
|
||||
INTEGRITY = ("Integrity", "Principles never break; reputation never fades", "🛡️")
|
||||
SELFLESSNESS = ("Selflessness", "Elevate others; personal success follows", "🤝")
|
||||
|
||||
# Science
|
||||
CURIOSITY = ("Curiosity", "Question always; truth never ends", "🌱")
|
||||
EMPIRICISM = ("Empiricism", "Hypothesize, measure; reality defines truth", "🔬")
|
||||
PRECISION = (
|
||||
"Precision",
|
||||
"Discipline in data; eliminate guesswork completely",
|
||||
"🎯",
|
||||
)
|
||||
VALIDATION = ("Validation", "Test assumptions hard; illusions crumble fast", "✅")
|
||||
OBJECTIVITY = ("Objectivity", "Ego out; results speak plainly", "🧊")
|
||||
REPEATABILITY = (
|
||||
"Repeatability",
|
||||
"Do it again; success repeats systematically",
|
||||
"🔄",
|
||||
)
|
||||
|
||||
# Design
|
||||
ACCESSIBILITY = (
|
||||
"Accessibility",
|
||||
"Open doors wide; adoption thrives naturally",
|
||||
"🌐",
|
||||
)
|
||||
BEAUTY = ("Beauty", "Form speaks louder; aesthetics lift utility", "🎨")
|
||||
CLARITY = ("Clarity", "Obvious is perfect; complexity hidden cleanly", "🔍")
|
||||
CONSISTENCY = ("Consistency", "Uniform patterns; predictable results always", "🎯")
|
||||
SIMPLICITY = ("Simplicity", "Cut ruthlessly; essential alone remains", "🪶")
|
||||
FLOW = ("Flow", "Remove friction; natural motion prevails", "🌊")
|
||||
|
||||
# Engineering
|
||||
BATTERIES_INCLUDED = (
|
||||
"Batteries Included",
|
||||
"Ready instantly; everything you need to start",
|
||||
"🔋",
|
||||
)
|
||||
CONCURRENCY = ("Concurrency", "Parallel flows; frictionless scale", "⚡")
|
||||
COMPOSABLE = ("Composable", "Modular magic; pieces multiply power", "🧩")
|
||||
INTEROPERABLE = (
|
||||
"Interoperable",
|
||||
"Integrate effortlessly; value compounds infinitely",
|
||||
"🔗",
|
||||
)
|
||||
ORTHOGONAL = ("Orthogonal", "Each tool exact; no overlap, no waste", "⚙️")
|
||||
SCALABLE = ("Scalable", "Growth limitless; obstacles removed at inception", "📈")
|
||||
|
||||
# Scale
|
||||
DISRUPTION = ("Disruption", "Reinvent boldly; transcend competition entirely", "💥")
|
||||
EXPERIMENTATION = ("Experimentation", "Test quickly; iterate endlessly", "🧪")
|
||||
EXPONENTIALITY = ("Exponentiality", "Compound constantly; incremental fades", "📈")
|
||||
VELOCITY = ("Velocity", "Ship fast; refine faster", "🚀")
|
||||
URGENCY = ("Urgency", "Act now; delays destroy opportunity", "⏱️")
|
||||
|
||||
# Wisdom
|
||||
ADAPTABILITY = (
|
||||
"Adaptability",
|
||||
"Pivot sharply; fluid response accelerates evolution",
|
||||
"🌊",
|
||||
)
|
||||
DECENTRALIZATION = (
|
||||
"Decentralization",
|
||||
"Distribute power; resilience born from autonomy",
|
||||
"🕸️",
|
||||
)
|
||||
FREEDOM = (
|
||||
"Freedom",
|
||||
"Democratize creativity; tools liberated, gatekeepers removed",
|
||||
"🗽",
|
||||
)
|
||||
LONGEVITY = (
|
||||
"Longevity",
|
||||
"Build timelessly; greatness endures beyond lifetimes",
|
||||
"⏳",
|
||||
)
|
||||
SECURITY = ("Security", "Encryption first; privacy non-negotiable", "🔐")
|
||||
ZEN = ("Zen", "Calm mastery; effortless excellence every moment", "☯️")
|
||||
|
||||
|
||||
class Hexagram:
|
||||
"""I Ching hexagram with interpretation."""
|
||||
|
||||
HEXAGRAMS = {
|
||||
"111111": (
|
||||
"乾 (Qián)",
|
||||
"Creative",
|
||||
"Initiating force, pure yang energy. Time for bold action.",
|
||||
),
|
||||
"000000": (
|
||||
"坤 (Kūn)",
|
||||
"Receptive",
|
||||
"Pure receptivity, yielding. Time to listen and adapt.",
|
||||
),
|
||||
"100010": (
|
||||
"屯 (Zhūn)",
|
||||
"Initial Difficulty",
|
||||
"Growing pains. Persevere through early challenges.",
|
||||
),
|
||||
"010001": (
|
||||
"蒙 (Méng)",
|
||||
"Youthful Folly",
|
||||
"Beginner's mind. Learn humbly, question assumptions.",
|
||||
),
|
||||
"111010": (
|
||||
"需 (Xū)",
|
||||
"Waiting",
|
||||
"Strategic patience. Prepare while waiting for the right moment.",
|
||||
),
|
||||
"010111": (
|
||||
"訟 (Sòng)",
|
||||
"Conflict",
|
||||
"Address conflicts directly but seek resolution, not victory.",
|
||||
),
|
||||
"010000": (
|
||||
"師 (Shī)",
|
||||
"Army",
|
||||
"Organize resources, build strong teams, lead by example.",
|
||||
),
|
||||
"000010": (
|
||||
"比 (Bǐ)",
|
||||
"Holding Together",
|
||||
"Unity and collaboration. Strengthen bonds.",
|
||||
),
|
||||
"111011": (
|
||||
"小畜 (Xiǎo Chù)",
|
||||
"Small Accumulation",
|
||||
"Small consistent improvements compound over time.",
|
||||
),
|
||||
"110111": (
|
||||
"履 (Lǚ)",
|
||||
"Treading",
|
||||
"Careful progress. Mind the details while moving forward.",
|
||||
),
|
||||
"111000": (
|
||||
"泰 (Tài)",
|
||||
"Peace",
|
||||
"Harmony achieved. Maintain balance while building.",
|
||||
),
|
||||
"000111": (
|
||||
"否 (Pǐ)",
|
||||
"Standstill",
|
||||
"Blockage present. Pause, reassess, find new paths.",
|
||||
),
|
||||
"101111": (
|
||||
"同人 (Tóng Rén)",
|
||||
"Fellowship",
|
||||
"Community strength. Build alliances and share knowledge.",
|
||||
),
|
||||
"111101": (
|
||||
"大有 (Dà Yǒu)",
|
||||
"Great Possession",
|
||||
"Abundance available. Share generously to multiply value.",
|
||||
),
|
||||
"001000": (
|
||||
"謙 (Qiān)",
|
||||
"Modesty",
|
||||
"Humble confidence. Let work speak for itself.",
|
||||
),
|
||||
"000100": (
|
||||
"豫 (Yù)",
|
||||
"Enthusiasm",
|
||||
"Infectious energy. Channel excitement into action.",
|
||||
),
|
||||
"100110": (
|
||||
"隨 (Suí)",
|
||||
"Following",
|
||||
"Adaptive leadership. Know when to lead and when to follow.",
|
||||
),
|
||||
"011001": (
|
||||
"蠱 (Gǔ)",
|
||||
"Work on Decay",
|
||||
"Fix technical debt. Address root causes.",
|
||||
),
|
||||
"110000": (
|
||||
"臨 (Lín)",
|
||||
"Approach",
|
||||
"Opportunity approaching. Prepare to receive it.",
|
||||
),
|
||||
"000011": (
|
||||
"觀 (Guān)",
|
||||
"Contemplation",
|
||||
"Step back for perspective. See the whole system.",
|
||||
),
|
||||
"100101": (
|
||||
"噬嗑 (Shì Kè)",
|
||||
"Biting Through",
|
||||
"Remove obstacles decisively. Clear blockages.",
|
||||
),
|
||||
"101001": ("賁 (Bì)", "Grace", "Polish and refine. Beauty enhances function."),
|
||||
"000001": (
|
||||
"剝 (Bō)",
|
||||
"Splitting Apart",
|
||||
"Decay phase. Let go of what's not working.",
|
||||
),
|
||||
"100000": (
|
||||
"復 (Fù)",
|
||||
"Return",
|
||||
"New cycle begins. Start fresh with lessons learned.",
|
||||
),
|
||||
"100111": (
|
||||
"無妄 (Wú Wàng)",
|
||||
"Innocence",
|
||||
"Act with pure intention. Avoid overthinking.",
|
||||
),
|
||||
"111001": (
|
||||
"大畜 (Dà Chù)",
|
||||
"Great Accumulation",
|
||||
"Build reserves. Invest in infrastructure.",
|
||||
),
|
||||
"100001": (
|
||||
"頤 (Yí)",
|
||||
"Nourishment",
|
||||
"Feed growth. Provide resources teams need.",
|
||||
),
|
||||
"011110": (
|
||||
"大過 (Dà Guò)",
|
||||
"Great Excess",
|
||||
"Extraordinary measures needed. Bold action required.",
|
||||
),
|
||||
"010010": (
|
||||
"坎 (Kǎn)",
|
||||
"Abysmal",
|
||||
"Navigate danger carefully. Trust your training.",
|
||||
),
|
||||
"101101": (
|
||||
"離 (Lí)",
|
||||
"Clinging Fire",
|
||||
"Clarity and vision. Illuminate the path forward.",
|
||||
),
|
||||
"001110": (
|
||||
"咸 (Xián)",
|
||||
"Influence",
|
||||
"Mutual attraction. Build on natural affinities.",
|
||||
),
|
||||
"011100": (
|
||||
"恆 (Héng)",
|
||||
"Duration",
|
||||
"Persistence pays. Maintain steady effort.",
|
||||
),
|
||||
"001111": ("遯 (Dùn)", "Retreat", "Strategic withdrawal. Regroup and refocus."),
|
||||
"111100": (
|
||||
"大壯 (Dà Zhuàng)",
|
||||
"Great Power",
|
||||
"Strength available. Use power responsibly.",
|
||||
),
|
||||
"000101": (
|
||||
"晉 (Jìn)",
|
||||
"Progress",
|
||||
"Advance steadily. Each step builds momentum.",
|
||||
),
|
||||
"101000": (
|
||||
"明夷 (Míng Yí)",
|
||||
"Darkening Light",
|
||||
"Work quietly. Keep brilliance hidden for now.",
|
||||
),
|
||||
"101011": (
|
||||
"家人 (Jiā Rén)",
|
||||
"Family",
|
||||
"Team harmony. Strengthen internal culture.",
|
||||
),
|
||||
"110101": (
|
||||
"睽 (Kuí)",
|
||||
"Opposition",
|
||||
"Creative tension. Find synthesis in differences.",
|
||||
),
|
||||
"001010": (
|
||||
"蹇 (Jiǎn)",
|
||||
"Obstruction",
|
||||
"Difficulty ahead. Find alternative routes.",
|
||||
),
|
||||
"010100": (
|
||||
"解 (Xiè)",
|
||||
"Deliverance",
|
||||
"Breakthrough achieved. Consolidate gains.",
|
||||
),
|
||||
"110001": ("損 (Sǔn)", "Decrease", "Simplify ruthlessly. Less is more."),
|
||||
"100011": ("益 (Yì)", "Increase", "Multiply value. Invest in growth."),
|
||||
"111110": (
|
||||
"夬 (Guài)",
|
||||
"Breakthrough",
|
||||
"Decisive moment. Act with conviction.",
|
||||
),
|
||||
"011111": (
|
||||
"姤 (Gòu)",
|
||||
"Coming to Meet",
|
||||
"Unexpected encounter. Stay alert to opportunity.",
|
||||
),
|
||||
"000110": (
|
||||
"萃 (Cuì)",
|
||||
"Gathering",
|
||||
"Convergence point. Bring elements together.",
|
||||
),
|
||||
"011000": (
|
||||
"升 (Shēng)",
|
||||
"Pushing Upward",
|
||||
"Gradual ascent. Build systematically.",
|
||||
),
|
||||
"010110": ("困 (Kùn)", "Exhaustion", "Resources depleted. Rest and recharge."),
|
||||
"011010": ("井 (Jǐng)", "The Well", "Deep resources. Draw from fundamentals."),
|
||||
"101110": (
|
||||
"革 (Gé)",
|
||||
"Revolution",
|
||||
"Transform completely. Embrace radical change.",
|
||||
),
|
||||
"011101": (
|
||||
"鼎 (Dǐng)",
|
||||
"The Cauldron",
|
||||
"Transformation vessel. Cook new solutions.",
|
||||
),
|
||||
"100100": (
|
||||
"震 (Zhèn)",
|
||||
"Thunder",
|
||||
"Shocking awakening. Respond to wake-up calls.",
|
||||
),
|
||||
"001001": (
|
||||
"艮 (Gèn)",
|
||||
"Mountain",
|
||||
"Stillness and stability. Find solid ground.",
|
||||
),
|
||||
"001011": (
|
||||
"漸 (Jiàn)",
|
||||
"Gradual Progress",
|
||||
"Step by step. Patient development.",
|
||||
),
|
||||
"110100": (
|
||||
"歸妹 (Guī Mèi)",
|
||||
"Marrying Maiden",
|
||||
"New partnerships. Align expectations.",
|
||||
),
|
||||
"101100": ("豐 (Fēng)", "Abundance", "Peak achievement. Prepare for cycles."),
|
||||
"001101": ("旅 (Lǚ)", "The Wanderer", "Explorer mindset. Learn from journey."),
|
||||
"011011": (
|
||||
"巽 (Xùn)",
|
||||
"Gentle Wind",
|
||||
"Subtle influence. Persistent gentle pressure.",
|
||||
),
|
||||
"110110": ("兌 (Duì)", "Joy", "Infectious happiness. Celebrate progress."),
|
||||
"010011": ("渙 (Huàn)", "Dispersion", "Break up rigidity. Dissolve barriers."),
|
||||
"110010": (
|
||||
"節 (Jié)",
|
||||
"Limitation",
|
||||
"Healthy constraints. Focus through limits.",
|
||||
),
|
||||
"110011": (
|
||||
"中孚 (Zhōng Fú)",
|
||||
"Inner Truth",
|
||||
"Authentic core. Build from truth.",
|
||||
),
|
||||
"001100": (
|
||||
"小過 (Xiǎo Guò)",
|
||||
"Small Excess",
|
||||
"Minor adjustments. Fine-tune carefully.",
|
||||
),
|
||||
"101010": (
|
||||
"既濟 (Jì Jì)",
|
||||
"After Completion",
|
||||
"Success achieved. Maintain vigilance.",
|
||||
),
|
||||
"010101": (
|
||||
"未濟 (Wèi Jì)",
|
||||
"Before Completion",
|
||||
"Almost there. Final push needed.",
|
||||
),
|
||||
}
|
||||
|
||||
def __init__(self, lines: str):
|
||||
self.lines = lines
|
||||
self.name, self.title, self.meaning = self.HEXAGRAMS.get(
|
||||
lines,
|
||||
("Unknown", "Mystery", "The pattern is unclear. Trust your intuition."),
|
||||
)
|
||||
|
||||
def get_changing_lines(self) -> List[int]:
|
||||
"""Identify which lines are changing (would be 6 or 9 in traditional I Ching)."""
|
||||
# For simplicity, randomly select 0-2 changing lines
|
||||
num_changes = random.choice([0, 1, 1, 2])
|
||||
if num_changes == 0:
|
||||
return []
|
||||
positions = list(range(6))
|
||||
return sorted(random.sample(positions, num_changes))
|
||||
|
||||
|
||||
class IChing:
|
||||
"""I Ching oracle for engineering guidance."""
|
||||
|
||||
def __init__(self):
|
||||
self.principles = list(HanzoPrinciple)
|
||||
|
||||
def cast_hexagram(self) -> Hexagram:
|
||||
"""Cast a hexagram using virtual coins."""
|
||||
lines = ""
|
||||
for _ in range(6):
|
||||
# Three coin tosses: heads=3, tails=2
|
||||
coins = sum(random.choice([2, 3]) for _ in range(3))
|
||||
# 6=old yin(changing 0), 7=young yang(1), 8=young yin(0), 9=old yang(changing 1)
|
||||
if coins in [6, 8]:
|
||||
lines += "0"
|
||||
else:
|
||||
lines += "1"
|
||||
return Hexagram(lines)
|
||||
|
||||
def select_principles(
|
||||
self, hexagram: Hexagram, challenge: str
|
||||
) -> List[HanzoPrinciple]:
|
||||
"""Select relevant Hanzo principles based on hexagram and challenge."""
|
||||
# Use hexagram pattern to deterministically but creatively select principles
|
||||
selected = []
|
||||
|
||||
# Primary principle based on hexagram pattern
|
||||
primary_index = sum(
|
||||
int(bit) * (2**i) for i, bit in enumerate(hexagram.lines)
|
||||
) % len(self.principles)
|
||||
selected.append(self.principles[primary_index])
|
||||
|
||||
# Supporting principles based on challenge keywords
|
||||
keywords = challenge.lower().split()
|
||||
keyword_matches = {
|
||||
"scale": [HanzoPrinciple.SCALABLE, HanzoPrinciple.EXPONENTIALITY],
|
||||
"speed": [HanzoPrinciple.VELOCITY, HanzoPrinciple.URGENCY],
|
||||
"quality": [HanzoPrinciple.PRECISION, HanzoPrinciple.VALIDATION],
|
||||
"team": [HanzoPrinciple.AUTONOMY, HanzoPrinciple.BALANCE],
|
||||
"design": [HanzoPrinciple.SIMPLICITY, HanzoPrinciple.BEAUTY],
|
||||
"bug": [HanzoPrinciple.EMPIRICISM, HanzoPrinciple.OBJECTIVITY],
|
||||
"refactor": [HanzoPrinciple.CLARITY, HanzoPrinciple.COMPOSABLE],
|
||||
"security": [HanzoPrinciple.SECURITY, HanzoPrinciple.INTEGRITY],
|
||||
"performance": [HanzoPrinciple.CONCURRENCY, HanzoPrinciple.ORTHOGONAL],
|
||||
"user": [HanzoPrinciple.CUSTOMER_OBSESSION, HanzoPrinciple.ACCESSIBILITY],
|
||||
}
|
||||
|
||||
for keyword, principles in keyword_matches.items():
|
||||
if keyword in keywords:
|
||||
selected.extend(principles)
|
||||
|
||||
# Add complementary principle based on changing lines
|
||||
changing_lines = hexagram.get_changing_lines()
|
||||
if changing_lines:
|
||||
complement_index = (primary_index + sum(changing_lines)) % len(
|
||||
self.principles
|
||||
)
|
||||
selected.append(self.principles[complement_index])
|
||||
|
||||
# Ensure uniqueness and limit to 3-5 principles
|
||||
seen = set()
|
||||
unique_selected = []
|
||||
for principle in selected:
|
||||
if principle not in seen:
|
||||
seen.add(principle)
|
||||
unique_selected.append(principle)
|
||||
|
||||
return unique_selected[:5]
|
||||
|
||||
def generate_guidance(
|
||||
self, hexagram: Hexagram, principles: List[HanzoPrinciple], challenge: str
|
||||
) -> str:
|
||||
"""Generate creative guidance combining I Ching wisdom and Hanzo principles."""
|
||||
guidance = f"☯️ I CHING GUIDANCE FOR ENGINEERING CHALLENGE ☯️\n\n"
|
||||
guidance += f"**Your Challenge:** {challenge}\n\n"
|
||||
|
||||
guidance += f"**Hexagram Cast:** {hexagram.name} - {hexagram.title}\n"
|
||||
guidance += f"**Pattern:** {''.join('━━━' if l == '1' else '━ ━' for l in hexagram.lines[::-1])}\n"
|
||||
guidance += f"**Ancient Wisdom:** {hexagram.meaning}\n\n"
|
||||
|
||||
guidance += "**Hanzo Principles to Apply:**\n\n"
|
||||
|
||||
for principle in principles:
|
||||
name, wisdom, emoji = principle.value
|
||||
guidance += f"{emoji} **{name}**\n"
|
||||
guidance += f" *{wisdom}*\n\n"
|
||||
|
||||
# Generate specific actionable advice
|
||||
guidance += "**Synthesized Approach:**\n\n"
|
||||
|
||||
# Hexagram-specific guidance
|
||||
if "Creative" in hexagram.title:
|
||||
guidance += "• This is a time for bold innovation. Don't hold back on ambitious ideas.\n"
|
||||
elif "Receptive" in hexagram.title:
|
||||
guidance += (
|
||||
"• Listen deeply to user needs and system constraints before acting.\n"
|
||||
)
|
||||
elif "Difficulty" in hexagram.title:
|
||||
guidance += (
|
||||
"• Challenges are teachers. Each obstacle reveals the path forward.\n"
|
||||
)
|
||||
elif "Waiting" in hexagram.title:
|
||||
guidance += "• Strategic patience required. Prepare thoroughly before implementation.\n"
|
||||
elif "Conflict" in hexagram.title:
|
||||
guidance += "• Technical disagreements? Seek data-driven resolution.\n"
|
||||
elif "Peace" in hexagram.title:
|
||||
guidance += (
|
||||
"• Harmony achieved. Now build sustainably on this foundation.\n"
|
||||
)
|
||||
|
||||
# Principle-specific actionable advice
|
||||
principle_actions = {
|
||||
HanzoPrinciple.SCALABLE: "• Design for 10x growth from day one. Remove scaling bottlenecks now.",
|
||||
HanzoPrinciple.VELOCITY: "• Ship an MVP today. Perfect is the enemy of shipped.",
|
||||
HanzoPrinciple.SIMPLICITY: "• Delete half your code. The best code is no code.",
|
||||
HanzoPrinciple.EMPIRICISM: "• Measure everything. Let data guide your decisions.",
|
||||
HanzoPrinciple.CUSTOMER_OBSESSION: "• Talk to users now. Their pain is your roadmap.",
|
||||
HanzoPrinciple.CONCURRENCY: "• Parallelize everything possible. Sequential is slow.",
|
||||
HanzoPrinciple.SECURITY: "• Security is not optional. Encrypt by default.",
|
||||
HanzoPrinciple.ZEN: "• Find calm in the chaos. Clear mind writes better code.",
|
||||
}
|
||||
|
||||
for principle in principles:
|
||||
if principle in principle_actions:
|
||||
guidance += principle_actions[principle] + "\n"
|
||||
|
||||
# Changing lines wisdom
|
||||
changing_lines = hexagram.get_changing_lines()
|
||||
if changing_lines:
|
||||
guidance += f"\n**Lines in Transition:** {', '.join(str(i + 1) for i in changing_lines)}\n"
|
||||
guidance += (
|
||||
"• Change is imminent in these areas. Prepare for transformation.\n"
|
||||
)
|
||||
|
||||
# Final synthesis
|
||||
guidance += "\n**The Way Forward:**\n"
|
||||
guidance += self._synthesize_action_plan(hexagram, principles, challenge)
|
||||
|
||||
guidance += "\n\n*Remember: The I Ching reveals patterns, not prescriptions. "
|
||||
guidance += "Let this wisdom guide your intuition as you craft your solution.*"
|
||||
|
||||
return guidance
|
||||
|
||||
def _synthesize_action_plan(
|
||||
self, hexagram: Hexagram, principles: List[HanzoPrinciple], challenge: str
|
||||
) -> str:
|
||||
"""Create a specific action plan based on the reading."""
|
||||
plan = ""
|
||||
|
||||
# Determine the nature of the challenge
|
||||
if any(word in challenge.lower() for word in ["bug", "error", "fix", "broken"]):
|
||||
plan += "1. **Diagnose systematically** - Use empirical debugging, not guesswork\n"
|
||||
plan += "2. **Fix root cause** - Address the source, not just symptoms\n"
|
||||
plan += "3. **Prevent recurrence** - Add tests and monitoring\n"
|
||||
elif any(
|
||||
word in challenge.lower() for word in ["scale", "performance", "slow"]
|
||||
):
|
||||
plan += "1. **Measure first** - Profile to find actual bottlenecks\n"
|
||||
plan += "2. **Parallelize** - Use concurrency where possible\n"
|
||||
plan += "3. **Simplify** - Remove complexity before optimizing\n"
|
||||
elif any(
|
||||
word in challenge.lower() for word in ["design", "architect", "structure"]
|
||||
):
|
||||
plan += "1. **Start simple** - MVP first, elaborate later\n"
|
||||
plan += "2. **Stay flexible** - Design for change\n"
|
||||
plan += "3. **Think holistically** - Consider entire system\n"
|
||||
elif any(
|
||||
word in challenge.lower() for word in ["team", "collaborate", "people"]
|
||||
):
|
||||
plan += "1. **Enable autonomy** - Trust your team\n"
|
||||
plan += "2. **Maintain balance** - Sustainable pace wins\n"
|
||||
plan += "3. **Share knowledge** - Elevate everyone\n"
|
||||
else:
|
||||
plan += "1. **Clarify intent** - What problem are you really solving?\n"
|
||||
plan += "2. **Start small** - Build incrementally\n"
|
||||
plan += "3. **Iterate rapidly** - Fast feedback loops\n"
|
||||
|
||||
return plan
|
||||
|
||||
|
||||
class IChingTool(BaseTool):
|
||||
"""Tool for applying I Ching wisdom to engineering challenges."""
|
||||
|
||||
name = "iching"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Apply I Ching wisdom and Hanzo principles to engineering challenges.
|
||||
|
||||
This tool casts an I Ching hexagram and selects relevant Hanzo principles
|
||||
to provide creative guidance for your engineering challenge.
|
||||
|
||||
Parameters:
|
||||
- challenge: Description of the engineering challenge or question
|
||||
|
||||
The oracle will:
|
||||
1. Cast a hexagram representing the current situation
|
||||
2. Select relevant Hanzo principles
|
||||
3. Synthesize actionable guidance
|
||||
4. Provide specific recommendations
|
||||
|
||||
Example:
|
||||
iching(
|
||||
challenge="How should I approach refactoring this legacy codebase?"
|
||||
)
|
||||
|
||||
Use this when you need:
|
||||
- Fresh perspective on a problem
|
||||
- Creative approach to challenges
|
||||
- Wisdom for difficult decisions
|
||||
- Alignment with Hanzo principles"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the I Ching tool."""
|
||||
super().__init__()
|
||||
self.oracle = IChing()
|
||||
|
||||
async def call(self, ctx: MCPContext, challenge: str) -> str:
|
||||
"""Cast I Ching and provide guidance."""
|
||||
# Cast hexagram
|
||||
hexagram = self.oracle.cast_hexagram()
|
||||
|
||||
# Select relevant principles
|
||||
principles = self.oracle.select_principles(hexagram, challenge)
|
||||
|
||||
# Generate guidance
|
||||
guidance = self.oracle.generate_guidance(hexagram, principles, challenge)
|
||||
|
||||
return guidance
|
||||
|
||||
def register(self, server: FastMCP) -> None:
|
||||
"""Register the tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@server.tool(name=self.name, description=self.description)
|
||||
async def iching(ctx: MCPContext, challenge: str) -> str:
|
||||
return await tool_self.call(ctx, challenge)
|
||||
@@ -1,243 +0,0 @@
|
||||
"""Network tool for dispatching work to agent networks.
|
||||
|
||||
This tool enables distributed AI workloads across local and remote agent networks,
|
||||
with support for both local-only execution (via hanzo-miner) and cloud fallback.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
from typing import (
|
||||
List,
|
||||
Unpack,
|
||||
Optional,
|
||||
Annotated,
|
||||
TypedDict,
|
||||
final,
|
||||
override,
|
||||
)
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
# Import hanzo cluster if available
|
||||
try:
|
||||
from hanzoai import cluster
|
||||
|
||||
CLUSTER_AVAILABLE = True
|
||||
except ImportError:
|
||||
CLUSTER_AVAILABLE = False
|
||||
|
||||
|
||||
class NetworkToolParams(TypedDict, total=False):
|
||||
"""Parameters for the network tool."""
|
||||
|
||||
task: str
|
||||
agents: Optional[List[str]]
|
||||
mode: Optional[str] # "local", "distributed", "hybrid"
|
||||
model: Optional[str]
|
||||
routing: Optional[str] # "sequential", "parallel", "consensus"
|
||||
require_local: Optional[bool]
|
||||
|
||||
|
||||
@final
|
||||
class NetworkTool(BaseTool):
|
||||
"""Dispatch work to agent networks for distributed AI processing.
|
||||
|
||||
Modes:
|
||||
- local: Use only local compute (via hanzo-cluster/miner)
|
||||
- distributed: Use available network resources
|
||||
- hybrid: Prefer local, fallback to cloud
|
||||
|
||||
This tool is the evolution of the swarm tool, providing:
|
||||
- True distributed execution across devices
|
||||
- Local-first privacy-preserving AI
|
||||
- Automatic routing and load balancing
|
||||
- Integration with hanzo-miner for compute contribution
|
||||
"""
|
||||
|
||||
name = "network"
|
||||
description = "Dispatch tasks to agent networks for distributed AI processing"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
default_mode: str = "hybrid",
|
||||
cluster_endpoint: str = None,
|
||||
):
|
||||
"""Initialize the network tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager
|
||||
default_mode: Default execution mode
|
||||
cluster_endpoint: Optional cluster endpoint
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.default_mode = default_mode
|
||||
self.cluster_endpoint = cluster_endpoint or os.environ.get(
|
||||
"HANZO_CLUSTER_ENDPOINT", "http://localhost:8000"
|
||||
)
|
||||
self._cluster = None
|
||||
|
||||
async def _ensure_cluster(self):
|
||||
"""Ensure we have a cluster connection."""
|
||||
if not CLUSTER_AVAILABLE:
|
||||
return None
|
||||
|
||||
if not self._cluster:
|
||||
try:
|
||||
# Try to connect to existing cluster
|
||||
self._cluster = cluster.HanzoCluster()
|
||||
# Check if cluster is running
|
||||
import httpx
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(f"{self.cluster_endpoint}/health")
|
||||
if response.status_code != 200:
|
||||
# Start local cluster if not running
|
||||
await self._cluster.start()
|
||||
except Exception:
|
||||
# Cluster not available
|
||||
self._cluster = None
|
||||
|
||||
return self._cluster
|
||||
|
||||
@override
|
||||
async def call(self, ctx: MCPContext, **params: Unpack[NetworkToolParams]) -> str:
|
||||
"""Execute a task on the agent network.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
task: Task description to execute
|
||||
agents: Optional list of specific agents to use
|
||||
mode: Execution mode (local/distributed/hybrid)
|
||||
model: Optional model preference
|
||||
routing: Routing strategy
|
||||
require_local: Require local-only execution
|
||||
|
||||
Returns:
|
||||
JSON string with results
|
||||
"""
|
||||
task = params.get("task", "")
|
||||
if not task:
|
||||
return json.dumps({"error": "Task description required", "success": False})
|
||||
|
||||
mode = params.get("mode", self.default_mode)
|
||||
agents_list = params.get("agents", [])
|
||||
model_pref = params.get("model")
|
||||
routing = params.get("routing", "sequential")
|
||||
require_local = params.get("require_local", False)
|
||||
|
||||
# Check if we should use local cluster
|
||||
use_local = mode in ["local", "hybrid"] or require_local
|
||||
|
||||
results = {
|
||||
"task": task,
|
||||
"mode": mode,
|
||||
"routing": routing,
|
||||
"agents_used": [],
|
||||
"results": [],
|
||||
"success": False,
|
||||
}
|
||||
|
||||
try:
|
||||
# Try local execution first if requested
|
||||
if use_local:
|
||||
cluster = await self._ensure_cluster()
|
||||
if cluster:
|
||||
try:
|
||||
# Execute on local cluster
|
||||
local_result = await cluster.inference(
|
||||
prompt=task,
|
||||
model=model_pref or "llama-3.2-3b",
|
||||
max_tokens=4000,
|
||||
)
|
||||
|
||||
results["agents_used"].append("local-cluster")
|
||||
results["results"].append(
|
||||
{
|
||||
"agent": "local-cluster",
|
||||
"response": local_result.get("choices", [{}])[0].get(
|
||||
"text", ""
|
||||
),
|
||||
"local": True,
|
||||
}
|
||||
)
|
||||
results["success"] = True
|
||||
|
||||
# If local succeeded and not hybrid, return
|
||||
if mode == "local" or (mode == "hybrid" and results["results"]):
|
||||
return json.dumps(results, indent=2)
|
||||
|
||||
except Exception as e:
|
||||
if require_local:
|
||||
results["error"] = f"Local execution failed: {str(e)}"
|
||||
return json.dumps(results, indent=2)
|
||||
|
||||
# Agent-based execution with concurrency
|
||||
if not results["success"] or mode in ["distributed", "hybrid"]:
|
||||
from hanzo_mcp.tools.agent.agent_tool import AgentTool
|
||||
agent = AgentTool(permission_manager=self.permission_manager, model=model_pref)
|
||||
concurrency = max(1, len(agents_list)) if agents_list else 5 if routing == "parallel" else 1
|
||||
agent_params = {"prompts": task, "concurrency": concurrency}
|
||||
agent_result = await agent.call(ctx, **agent_params)
|
||||
# Wrap agent_result as a simple result list
|
||||
results["agents_used"].append("agent")
|
||||
results["results"].append({"agent": "agent", "response": agent_result})
|
||||
results["success"] = True
|
||||
|
||||
except Exception as e:
|
||||
results["error"] = str(e)
|
||||
|
||||
return json.dumps(results, indent=2)
|
||||
|
||||
def register(self, server: FastMCP):
|
||||
"""Register the network tool with the server.
|
||||
|
||||
Args:
|
||||
server: FastMCP server instance
|
||||
"""
|
||||
tool = self
|
||||
|
||||
@server.tool(name=tool.name, description=tool.description)
|
||||
async def network_handler(
|
||||
ctx: MCPContext,
|
||||
task: Annotated[str, Field(description="Task to execute on the network")],
|
||||
agents: Annotated[
|
||||
Optional[List[str]], Field(description="Specific agents to use")
|
||||
] = None,
|
||||
mode: Annotated[
|
||||
Optional[str],
|
||||
Field(description="Execution mode: local, distributed, or hybrid"),
|
||||
] = None,
|
||||
model: Annotated[
|
||||
Optional[str], Field(description="Model preference")
|
||||
] = None,
|
||||
routing: Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Routing strategy: sequential, parallel, or consensus"
|
||||
),
|
||||
] = None,
|
||||
require_local: Annotated[
|
||||
Optional[bool], Field(description="Require local-only execution")
|
||||
] = None,
|
||||
) -> str:
|
||||
"""Dispatch work to agent networks."""
|
||||
params = NetworkToolParams(
|
||||
task=task,
|
||||
agents=agents,
|
||||
mode=mode,
|
||||
model=model,
|
||||
routing=routing,
|
||||
require_local=require_local,
|
||||
)
|
||||
return await tool.call(ctx, **params)
|
||||
|
||||
return tool
|
||||
|
||||
|
||||
# Remove swarm compatibility tool; swarm is an alias of agent with concurrency
|
||||
@@ -1,183 +0,0 @@
|
||||
"""Prompt generation utilities for agent tool.
|
||||
|
||||
This module provides functions for generating effective prompts for sub-agents,
|
||||
including filtering tools based on permissions and formatting system instructions.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
|
||||
def get_allowed_agent_tools(
|
||||
tools: list[BaseTool],
|
||||
permission_manager: PermissionManager,
|
||||
) -> list[BaseTool]:
|
||||
"""Filter tools available to the agent based on permissions.
|
||||
|
||||
Args:
|
||||
tools: List of available tools
|
||||
permission_manager: Permission manager for checking tool access
|
||||
|
||||
Returns:
|
||||
Filtered list of tools available to the agent
|
||||
"""
|
||||
# Get all tools except for the agent tool itself (avoid recursion)
|
||||
filtered_tools = [tool for tool in tools if tool.name != "agent"]
|
||||
|
||||
return filtered_tools
|
||||
|
||||
|
||||
def get_system_prompt(
|
||||
tools: list[BaseTool],
|
||||
permission_manager: PermissionManager,
|
||||
) -> str:
|
||||
"""Generate system prompt for the sub-agent.
|
||||
|
||||
Args:
|
||||
tools: List of available tools
|
||||
permission_manager: Permission manager for checking tool access
|
||||
|
||||
Returns:
|
||||
System prompt for the sub-agent
|
||||
"""
|
||||
# Get filtered tools
|
||||
filtered_tools = get_allowed_agent_tools(tools, permission_manager)
|
||||
|
||||
# Extract tool names for display
|
||||
tool_names = ", ".join(f"`{tool.name}`" for tool in filtered_tools)
|
||||
|
||||
# Base system prompt - agents always have edit tools
|
||||
system_prompt = f"""You are a Claude sub-agent with access to these tools: {tool_names}.
|
||||
|
||||
CAPABILITIES:
|
||||
1. You have FULL read and write access - you can create, edit, and modify files
|
||||
2. You can ask clarifying questions if needed - your response goes to the coordinating agent
|
||||
3. You work as part of a team of specialized agents
|
||||
4. Other agents may be available via MCP tools (look for tools named after agents)
|
||||
5. When relevant, share file names and code snippets
|
||||
6. Any file paths you return MUST be absolute. DO NOT use relative paths.
|
||||
7. You can only work with the absolute paths provided in your task prompt.
|
||||
|
||||
CLARIFICATION:
|
||||
- You can request clarification ONCE per task using the request_clarification tool
|
||||
- Use this when instructions are ambiguous or you need additional context
|
||||
- Types: AMBIGUOUS_INSTRUCTION, MISSING_CONTEXT, MULTIPLE_OPTIONS, CONFIRMATION_NEEDED, ADDITIONAL_INFO
|
||||
- The main loop will provide automated guidance based on context
|
||||
|
||||
CRITICAL REVIEW (Devil's Advocate):
|
||||
- Use the critic tool to get harsh, challenging feedback that attacks assumptions
|
||||
- The critic will find flaws and push for improvements aggressively
|
||||
- You can request up to 2 critic reviews per task
|
||||
- Review types: CODE_QUALITY, CORRECTNESS, PERFORMANCE, SECURITY, COMPLETENESS, BEST_PRACTICES, GENERAL
|
||||
- Use when you need someone to find what's wrong with your approach
|
||||
|
||||
BALANCED REVIEW:
|
||||
- Use the review tool for constructive, balanced code review
|
||||
- Provides objective assessment without predetermined bias
|
||||
- You can request up to 3 reviews per task
|
||||
- Focus areas: GENERAL, FUNCTIONALITY, READABILITY, MAINTAINABILITY, TESTING, DOCUMENTATION, ARCHITECTURE
|
||||
- Use for regular code review and feedback
|
||||
|
||||
CREATIVE GUIDANCE:
|
||||
- Use the iching tool when you need creative problem-solving approaches
|
||||
- Combines ancient I Ching wisdom with Hanzo engineering principles
|
||||
- Provides unique perspectives and actionable guidance
|
||||
- Use when stuck, need fresh ideas, or want philosophical alignment
|
||||
|
||||
EDITING GUIDELINES:
|
||||
- ALWAYS read the file first before attempting any edits
|
||||
- For edit tool: The old_string must match EXACTLY including all whitespace, tabs, and newlines
|
||||
- When copying text from read output, be careful with line numbers and indentation
|
||||
- If an edit fails due to whitespace mismatch, try reading the specific lines again
|
||||
- Prefer multi_edit when making multiple changes to the same file
|
||||
- Test your edits by verifying the exact string exists in the file first
|
||||
|
||||
COLLABORATION:
|
||||
- If you see MCP tools named after other agents, you can communicate with them
|
||||
- Use agent MCP tools to delegate specialized tasks or get expert opinions
|
||||
- Share context when communicating with other agents
|
||||
|
||||
RESPONSE FORMAT:
|
||||
- Begin with a summary of what you did or found
|
||||
- If you have questions, ask them clearly
|
||||
- Include details of any edits performed
|
||||
- Report any errors encountered
|
||||
- End with clear conclusions or next steps
|
||||
"""
|
||||
|
||||
return system_prompt
|
||||
|
||||
|
||||
def get_default_model(model_override: str | None = None) -> str:
|
||||
"""Get the default model for agent execution.
|
||||
|
||||
Args:
|
||||
model_override: Optional model override string in LiteLLM format (e.g., "openai/gpt-4o")
|
||||
|
||||
Returns:
|
||||
Model identifier string with provider prefix
|
||||
"""
|
||||
# Use model override if provided
|
||||
if model_override:
|
||||
# If in testing mode and using a test model, return as-is
|
||||
if model_override.startswith("test-model") or "TEST_MODE" in os.environ:
|
||||
return model_override
|
||||
|
||||
# If the model already has a provider prefix, return as-is
|
||||
if "/" in model_override:
|
||||
return model_override
|
||||
|
||||
# Otherwise, add the default provider prefix
|
||||
provider = os.environ.get("AGENT_PROVIDER", "openai")
|
||||
return f"{provider}/{model_override}"
|
||||
|
||||
# Fall back to environment variables
|
||||
# Default to Sonnet for cost efficiency
|
||||
model = os.environ.get("AGENT_MODEL", "claude-3-5-sonnet-20241022")
|
||||
|
||||
# Special cases for tests
|
||||
if (
|
||||
model.startswith("test-model")
|
||||
or "TEST_MODE" in os.environ
|
||||
and model == "claude-3-5-sonnet-20241022"
|
||||
):
|
||||
return model
|
||||
|
||||
provider = os.environ.get("AGENT_PROVIDER", "anthropic")
|
||||
|
||||
# Only add provider prefix if it's not already in the model name
|
||||
if "/" not in model and provider != "anthropic":
|
||||
return f"{provider}/{model}"
|
||||
elif "/" not in model and provider == "anthropic":
|
||||
return f"anthropic/{model}"
|
||||
elif "/" not in model:
|
||||
return f"openai/{model}"
|
||||
else:
|
||||
# Model already has a provider prefix
|
||||
return model
|
||||
|
||||
|
||||
def get_model_parameters(max_tokens: int | None = None) -> dict[str, Any]:
|
||||
"""Get model parameters from environment variables.
|
||||
|
||||
Args:
|
||||
max_tokens: Optional maximum tokens parameter override
|
||||
|
||||
Returns:
|
||||
Dictionary of model parameters
|
||||
"""
|
||||
params = {
|
||||
"temperature": float(os.environ.get("AGENT_TEMPERATURE", "0.7")),
|
||||
"timeout": int(os.environ.get("AGENT_API_TIMEOUT", "60")),
|
||||
}
|
||||
|
||||
# Add max_tokens if provided or if set in environment variable
|
||||
if max_tokens is not None:
|
||||
params["max_tokens"] = max_tokens
|
||||
elif os.environ.get("AGENT_MAX_TOKENS"):
|
||||
params["max_tokens"] = int(os.environ.get("AGENT_MAX_TOKENS", "1000"))
|
||||
|
||||
return params
|
||||
@@ -1,436 +0,0 @@
|
||||
"""Review tool for agents to request balanced code review from main loop."""
|
||||
|
||||
from enum import Enum
|
||||
from typing import List, Optional, override
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
|
||||
|
||||
class ReviewFocus(Enum):
|
||||
"""Types of review focus areas."""
|
||||
|
||||
GENERAL = "general"
|
||||
FUNCTIONALITY = "functionality"
|
||||
READABILITY = "readability"
|
||||
MAINTAINABILITY = "maintainability"
|
||||
TESTING = "testing"
|
||||
DOCUMENTATION = "documentation"
|
||||
ARCHITECTURE = "architecture"
|
||||
|
||||
|
||||
class ReviewTool(BaseTool):
|
||||
"""Tool for agents to request balanced code review from the main loop."""
|
||||
|
||||
name = "review"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Request a balanced, constructive code review from the main loop.
|
||||
|
||||
Unlike the critic tool (which plays devil's advocate), this provides:
|
||||
- Objective assessment of code quality
|
||||
- Recognition of what's done well
|
||||
- Constructive suggestions for improvement
|
||||
- Focus on practical concerns
|
||||
- No predetermined bias or harsh judgment
|
||||
|
||||
Parameters:
|
||||
- focus: Review focus area (GENERAL, FUNCTIONALITY, READABILITY, MAINTAINABILITY, TESTING, DOCUMENTATION, ARCHITECTURE)
|
||||
- work_description: Clear description of what you've implemented
|
||||
- code_snippets: Optional code snippets to review (as a list of strings)
|
||||
- file_paths: Optional list of file paths you've modified
|
||||
- context: Optional additional context about the implementation
|
||||
|
||||
The review will be balanced, highlighting both strengths and areas for improvement.
|
||||
|
||||
Example:
|
||||
review(
|
||||
focus="FUNCTIONALITY",
|
||||
work_description="Implemented auto-import feature for Go files",
|
||||
code_snippets=["func AddImport(file string, importPath string) error { ... }"],
|
||||
file_paths=["/path/to/import_handler.go"],
|
||||
context="This will be used to automatically fix missing imports in Go files"
|
||||
)"""
|
||||
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
focus: str,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]] = None,
|
||||
file_paths: Optional[List[str]] = None,
|
||||
context: Optional[str] = None,
|
||||
) -> str:
|
||||
"""This is a placeholder - actual implementation happens in AgentTool."""
|
||||
# This tool is handled specially in the agent execution
|
||||
return f"Review requested for: {work_description}"
|
||||
|
||||
def register(self, server: FastMCP) -> None:
|
||||
"""Register the tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@server.tool(name=self.name, description=self.description)
|
||||
async def review(
|
||||
ctx: MCPContext,
|
||||
focus: str,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]] = None,
|
||||
file_paths: Optional[List[str]] = None,
|
||||
context: Optional[str] = None,
|
||||
) -> str:
|
||||
return await tool_self.call(
|
||||
ctx, focus, work_description, code_snippets, file_paths, context
|
||||
)
|
||||
|
||||
|
||||
class BalancedReviewer:
|
||||
"""Provides balanced, constructive code reviews."""
|
||||
|
||||
def __init__(self):
|
||||
self.review_handlers = {
|
||||
ReviewFocus.GENERAL: self._review_general,
|
||||
ReviewFocus.FUNCTIONALITY: self._review_functionality,
|
||||
ReviewFocus.READABILITY: self._review_readability,
|
||||
ReviewFocus.MAINTAINABILITY: self._review_maintainability,
|
||||
ReviewFocus.TESTING: self._review_testing,
|
||||
ReviewFocus.DOCUMENTATION: self._review_documentation,
|
||||
ReviewFocus.ARCHITECTURE: self._review_architecture,
|
||||
}
|
||||
|
||||
def review(
|
||||
self,
|
||||
focus: ReviewFocus,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]] = None,
|
||||
file_paths: Optional[List[str]] = None,
|
||||
context: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Perform a balanced code review."""
|
||||
review_func = self.review_handlers.get(focus, self._review_general)
|
||||
return review_func(work_description, code_snippets, file_paths, context)
|
||||
|
||||
def _review_general(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
context: Optional[str],
|
||||
) -> str:
|
||||
"""Provide a general balanced review."""
|
||||
response = "📋 GENERAL CODE REVIEW:\n\n"
|
||||
response += f"**Work Reviewed:** {work_description}\n\n"
|
||||
|
||||
# Positive observations
|
||||
response += "**Positive Aspects:**\n"
|
||||
if "fix" in work_description.lower():
|
||||
response += "✓ Addressing identified issues proactively\n"
|
||||
if "implement" in work_description.lower():
|
||||
response += "✓ Adding new functionality to enhance the system\n"
|
||||
if code_snippets:
|
||||
response += "✓ Code structure appears organized\n"
|
||||
if file_paths and len(file_paths) == 1:
|
||||
response += "✓ Focused changes in a single file (good for reviewability)\n"
|
||||
elif file_paths and len(file_paths) > 1:
|
||||
response += "✓ Comprehensive approach across multiple files\n"
|
||||
|
||||
# Constructive suggestions
|
||||
response += "\n**Suggestions for Consideration:**\n"
|
||||
response += "• Ensure all edge cases are handled appropriately\n"
|
||||
response += "• Consider adding unit tests if not already present\n"
|
||||
response += "• Verify the changes integrate well with existing code\n"
|
||||
|
||||
if context:
|
||||
response += f"\n**Context Consideration:**\n{context}\n"
|
||||
response += "→ This context helps understand the implementation choices.\n"
|
||||
|
||||
# Summary
|
||||
response += "\n**Summary:**\n"
|
||||
response += "The implementation appears sound. Consider the suggestions above to further strengthen the code."
|
||||
|
||||
return response
|
||||
|
||||
def _review_functionality(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
context: Optional[str],
|
||||
) -> str:
|
||||
"""Review functionality aspects."""
|
||||
response = "📋 FUNCTIONALITY REVIEW:\n\n"
|
||||
response += f"**Implementation:** {work_description}\n\n"
|
||||
|
||||
response += "**Functional Assessment:**\n"
|
||||
|
||||
# Analyze code snippets if provided
|
||||
if code_snippets:
|
||||
for i, snippet in enumerate(code_snippets, 1):
|
||||
response += f"\nCode Snippet {i}:\n"
|
||||
|
||||
# Check for function definitions
|
||||
if "func " in snippet or "def " in snippet or "function " in snippet:
|
||||
response += "✓ Function definition looks properly structured\n"
|
||||
|
||||
# Check for error handling
|
||||
if "error" in snippet or "err" in snippet or "try" in snippet:
|
||||
response += "✓ Error handling is present\n"
|
||||
elif "return" in snippet:
|
||||
response += "• Consider adding error handling if applicable\n"
|
||||
|
||||
# Check for input validation
|
||||
if "if " in snippet or "check" in snippet.lower():
|
||||
response += "✓ Input validation appears to be present\n"
|
||||
|
||||
response += "\n**Functional Considerations:**\n"
|
||||
response += "• Does the implementation handle all expected inputs?\n"
|
||||
response += "• Are return values meaningful and consistent?\n"
|
||||
response += "• Is the functionality easily testable?\n"
|
||||
response += "• Does it integrate well with existing features?\n"
|
||||
|
||||
response += "\n**Overall:** The functionality appears to meet the described requirements."
|
||||
|
||||
return response
|
||||
|
||||
def _review_readability(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
context: Optional[str],
|
||||
) -> str:
|
||||
"""Review code readability."""
|
||||
response = "📋 READABILITY REVIEW:\n\n"
|
||||
|
||||
response += "**Readability Factors:**\n"
|
||||
|
||||
if code_snippets:
|
||||
total_lines = sum(snippet.count("\n") + 1 for snippet in code_snippets)
|
||||
avg_line_length = sum(
|
||||
len(line) for snippet in code_snippets for line in snippet.split("\n")
|
||||
) / max(total_lines, 1)
|
||||
|
||||
if avg_line_length < 80:
|
||||
response += "✓ Line lengths are reasonable\n"
|
||||
else:
|
||||
response += (
|
||||
"• Some lines might be too long, consider breaking them up\n"
|
||||
)
|
||||
|
||||
# Check naming
|
||||
has_good_names = any(
|
||||
any(
|
||||
word in snippet
|
||||
for word in ["Add", "Get", "Set", "Create", "Update", "Delete"]
|
||||
)
|
||||
for snippet in code_snippets
|
||||
)
|
||||
if has_good_names:
|
||||
response += "✓ Function/method names appear descriptive\n"
|
||||
|
||||
response += "\n**Readability Suggestions:**\n"
|
||||
response += "• Use meaningful variable and function names\n"
|
||||
response += "• Keep functions focused on a single responsibility\n"
|
||||
response += "• Add comments for complex logic sections\n"
|
||||
response += "• Maintain consistent indentation and formatting\n"
|
||||
|
||||
response += "\n**Overall:** Code readability appears acceptable with room for minor improvements."
|
||||
|
||||
return response
|
||||
|
||||
def _review_maintainability(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
context: Optional[str],
|
||||
) -> str:
|
||||
"""Review maintainability aspects."""
|
||||
response = "📋 MAINTAINABILITY REVIEW:\n\n"
|
||||
|
||||
response += "**Maintainability Factors:**\n"
|
||||
|
||||
# Check file organization
|
||||
if file_paths:
|
||||
if len(file_paths) == 1:
|
||||
response += "✓ Changes are localized to a single file\n"
|
||||
else:
|
||||
response += "✓ Changes are logically distributed across files\n"
|
||||
|
||||
# Check for modularity in code
|
||||
if code_snippets:
|
||||
function_count = sum(
|
||||
snippet.count("func ")
|
||||
+ snippet.count("def ")
|
||||
+ snippet.count("function ")
|
||||
for snippet in code_snippets
|
||||
)
|
||||
if function_count > 0:
|
||||
response += "✓ Code is broken into functions/methods\n"
|
||||
|
||||
response += "\n**Maintainability Considerations:**\n"
|
||||
response += "• Is the code modular and reusable?\n"
|
||||
response += "• Are dependencies clearly defined?\n"
|
||||
response += "• Will future developers understand the intent?\n"
|
||||
response += "• Is the code structured to allow easy updates?\n"
|
||||
|
||||
response += "\n**Recommendations:**\n"
|
||||
response += "• Consider extracting common patterns into utilities\n"
|
||||
response += "• Ensure consistent patterns across the codebase\n"
|
||||
response += "• Document any non-obvious design decisions\n"
|
||||
|
||||
return response
|
||||
|
||||
def _review_testing(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
context: Optional[str],
|
||||
) -> str:
|
||||
"""Review testing aspects."""
|
||||
response = "📋 TESTING REVIEW:\n\n"
|
||||
|
||||
has_test_files = any("test" in str(path).lower() for path in (file_paths or []))
|
||||
|
||||
if has_test_files:
|
||||
response += "✓ Test files are included with the changes\n\n"
|
||||
else:
|
||||
response += "⚠️ No test files detected in the changes\n\n"
|
||||
|
||||
response += "**Testing Checklist:**\n"
|
||||
response += "□ Unit tests for new functions\n"
|
||||
response += "□ Integration tests for feature interactions\n"
|
||||
response += "□ Edge case coverage\n"
|
||||
response += "□ Error condition testing\n"
|
||||
response += "□ Performance tests (if applicable)\n"
|
||||
|
||||
response += "\n**Testing Recommendations:**\n"
|
||||
response += "• Write tests that document expected behavior\n"
|
||||
response += "• Include both positive and negative test cases\n"
|
||||
response += "• Ensure tests are maintainable and clear\n"
|
||||
response += "• Aim for good coverage of critical paths\n"
|
||||
|
||||
if not has_test_files:
|
||||
response += "\n💡 Consider adding tests to ensure reliability and prevent regressions."
|
||||
|
||||
return response
|
||||
|
||||
def _review_documentation(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
context: Optional[str],
|
||||
) -> str:
|
||||
"""Review documentation aspects."""
|
||||
response = "📋 DOCUMENTATION REVIEW:\n\n"
|
||||
|
||||
# Check for documentation in code
|
||||
has_comments = False
|
||||
if code_snippets:
|
||||
has_comments = any(
|
||||
"//" in snippet or "/*" in snippet or "#" in snippet or '"""' in snippet
|
||||
for snippet in code_snippets
|
||||
)
|
||||
|
||||
if has_comments:
|
||||
response += "✓ Code includes some documentation\n"
|
||||
else:
|
||||
response += "• Consider adding documentation comments\n"
|
||||
|
||||
response += "\n**Documentation Guidelines:**\n"
|
||||
response += "• Document the 'why' not just the 'what'\n"
|
||||
response += "• Include examples for complex functions\n"
|
||||
response += "• Document any assumptions or limitations\n"
|
||||
response += "• Keep documentation up-to-date with code changes\n"
|
||||
|
||||
response += "\n**Recommended Documentation:**\n"
|
||||
response += "• Function/method purpose and parameters\n"
|
||||
response += "• Complex algorithm explanations\n"
|
||||
response += "• API usage examples\n"
|
||||
response += "• Configuration requirements\n"
|
||||
|
||||
return response
|
||||
|
||||
def _review_architecture(
|
||||
self,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]],
|
||||
file_paths: Optional[List[str]],
|
||||
context: Optional[str],
|
||||
) -> str:
|
||||
"""Review architectural aspects."""
|
||||
response = "📋 ARCHITECTURE REVIEW:\n\n"
|
||||
|
||||
response += "**Architectural Considerations:**\n"
|
||||
|
||||
# Analyze file structure
|
||||
if file_paths:
|
||||
# Check for separation of concerns
|
||||
has_separation = (
|
||||
len(set(str(p).split("/")[-2] for p in file_paths if "/" in str(p))) > 1
|
||||
)
|
||||
if has_separation:
|
||||
response += "✓ Changes span multiple modules (good separation)\n"
|
||||
else:
|
||||
response += "✓ Changes are cohesive within a module\n"
|
||||
|
||||
response += "\n**Architectural Principles:**\n"
|
||||
response += "• Single Responsibility - Each component has one clear purpose\n"
|
||||
response += "• Open/Closed - Open for extension, closed for modification\n"
|
||||
response += "• Dependency Inversion - Depend on abstractions, not concretions\n"
|
||||
response += "• Interface Segregation - Keep interfaces focused and minimal\n"
|
||||
|
||||
response += "\n**Questions to Consider:**\n"
|
||||
response += "• Does this fit well with the existing architecture?\n"
|
||||
response += "• Are the right abstractions in place?\n"
|
||||
response += "• Is the coupling between components appropriate?\n"
|
||||
response += "• Will this scale as requirements grow?\n"
|
||||
|
||||
if context:
|
||||
response += f"\n**Context Impact:**\n{context}\n"
|
||||
response += "→ Ensure the architectural choices align with this context.\n"
|
||||
|
||||
return response
|
||||
|
||||
|
||||
class ReviewProtocol:
|
||||
"""Protocol for review interactions."""
|
||||
|
||||
def __init__(self):
|
||||
self.reviewer = BalancedReviewer()
|
||||
self.review_count = 0
|
||||
self.max_reviews = 3 # Allow up to 3 reviews per task
|
||||
|
||||
def request_review(
|
||||
self,
|
||||
focus: str,
|
||||
work_description: str,
|
||||
code_snippets: Optional[List[str]] = None,
|
||||
file_paths: Optional[List[str]] = None,
|
||||
context: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Request a balanced review."""
|
||||
if self.review_count >= self.max_reviews:
|
||||
return "📋 Review limit reached. You've received comprehensive feedback - time to finalize your implementation."
|
||||
|
||||
self.review_count += 1
|
||||
|
||||
try:
|
||||
focus_enum = ReviewFocus[focus.upper()]
|
||||
except KeyError:
|
||||
focus_enum = ReviewFocus.GENERAL
|
||||
|
||||
review = self.reviewer.review(
|
||||
focus_enum, work_description, code_snippets, file_paths, context
|
||||
)
|
||||
|
||||
header = f"Review {self.review_count}/{self.max_reviews} (Focus: {focus_enum.value}):\n\n"
|
||||
footer = "\n\n💡 This is a balanced review - consider both strengths and suggestions."
|
||||
|
||||
return header + review + footer
|
||||
@@ -1,87 +0,0 @@
|
||||
"""Swarm tool as an alias to Network tool for backward compatibility.
|
||||
|
||||
This module makes swarm an alias to the network tool, as network is the
|
||||
evolution of swarm with better distributed execution capabilities.
|
||||
"""
|
||||
|
||||
from hanzo_mcp.tools.agent.network_tool import NetworkTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
|
||||
class SwarmTool(NetworkTool):
|
||||
"""Swarm tool - alias to Network tool for backward compatibility.
|
||||
|
||||
The swarm tool is now an alias to the network tool, which provides
|
||||
all the same functionality plus additional distributed execution modes.
|
||||
Use 'network' for new code, 'swarm' is maintained for compatibility.
|
||||
"""
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "swarm"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Execute a network of AI agents (alias to 'network' tool).
|
||||
|
||||
The 'swarm' tool is now an alias to the 'network' tool for backward compatibility.
|
||||
All swarm functionality is available through network, which additionally provides:
|
||||
|
||||
- Local-first execution with privacy preservation
|
||||
- Distributed compute across devices
|
||||
- Hybrid mode with cloud fallback
|
||||
- Integration with hanzo-network for MCP-connected agents
|
||||
|
||||
Examples:
|
||||
```python
|
||||
# These are equivalent:
|
||||
swarm(task="Analyze code", agents=["analyzer", "reviewer"])
|
||||
network(task="Analyze code", agents=["analyzer", "reviewer"])
|
||||
|
||||
# Network adds new modes:
|
||||
network(task="Process data", mode="local") # Privacy-first
|
||||
network(task="Large analysis", mode="distributed") # Scale out
|
||||
```
|
||||
|
||||
For new code, prefer using 'network' directly."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
default_mode: str = "hybrid",
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize swarm as an alias to network.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager
|
||||
default_mode: Default execution mode (hybrid/local/distributed)
|
||||
**kwargs: Additional arguments passed to NetworkTool
|
||||
"""
|
||||
# Just pass through to NetworkTool
|
||||
super().__init__(
|
||||
permission_manager=permission_manager, default_mode=default_mode, **kwargs
|
||||
)
|
||||
|
||||
async def call(self, **kwargs) -> str:
|
||||
"""Execute swarm via network tool.
|
||||
|
||||
All parameters are passed through to the network tool.
|
||||
"""
|
||||
# For backward compatibility, rename some parameters if needed
|
||||
if "config" in kwargs and "agents" not in kwargs:
|
||||
# Old swarm used 'config' for agent definitions
|
||||
config = kwargs.pop("config")
|
||||
if isinstance(config, dict) and "agents" in config:
|
||||
kwargs["agents"] = config["agents"]
|
||||
if "topology" in config:
|
||||
kwargs["routing"] = config["topology"]
|
||||
|
||||
# Pass through to network tool
|
||||
return await super().call(**kwargs)
|
||||
|
||||
|
||||
# For backward compatibility exports
|
||||
__all__ = ["SwarmTool"]
|
||||
@@ -1,718 +0,0 @@
|
||||
"""Swarm tool implementation using hanzo-agents SDK.
|
||||
|
||||
This module implements the SwarmTool that leverages the hanzo-agents SDK
|
||||
for sophisticated multi-agent orchestration with flexible network topologies.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
List,
|
||||
Unpack,
|
||||
Optional,
|
||||
TypedDict,
|
||||
final,
|
||||
override,
|
||||
)
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
# Import hanzo-agents SDK with fallback
|
||||
try:
|
||||
from hanzo_agents import (
|
||||
Tool,
|
||||
Agent,
|
||||
State,
|
||||
Router,
|
||||
History,
|
||||
Network,
|
||||
ToolCall,
|
||||
ModelRegistry,
|
||||
InferenceResult,
|
||||
)
|
||||
|
||||
HANZO_AGENTS_AVAILABLE = True
|
||||
except ImportError:
|
||||
# Define minimal stubs if hanzo-agents is not available
|
||||
HANZO_AGENTS_AVAILABLE = False
|
||||
|
||||
class Agent:
|
||||
pass
|
||||
|
||||
class State:
|
||||
pass
|
||||
|
||||
class Network:
|
||||
pass
|
||||
|
||||
class Tool:
|
||||
pass
|
||||
|
||||
class History:
|
||||
pass
|
||||
|
||||
class ModelRegistry:
|
||||
pass
|
||||
|
||||
class InferenceResult:
|
||||
pass
|
||||
|
||||
class ToolCall:
|
||||
pass
|
||||
|
||||
class Router:
|
||||
pass
|
||||
|
||||
|
||||
# Import optional components with fallbacks
|
||||
try:
|
||||
from hanzo_agents import LLMRouter, HybridRouter, DeterministicRouter
|
||||
except ImportError:
|
||||
try:
|
||||
# Try core module import
|
||||
from hanzo_agents.core.router import (
|
||||
LLMRouter,
|
||||
HybridRouter,
|
||||
DeterministicRouter,
|
||||
)
|
||||
except ImportError:
|
||||
# Define stubs if not available
|
||||
class DeterministicRouter:
|
||||
pass
|
||||
|
||||
class LLMRouter:
|
||||
pass
|
||||
|
||||
class HybridRouter:
|
||||
pass
|
||||
|
||||
|
||||
try:
|
||||
from hanzo_agents import create_memory_kv, create_memory_vector
|
||||
except ImportError:
|
||||
try:
|
||||
# Try core module import
|
||||
from hanzo_agents.core.memory import create_memory_kv, create_memory_vector
|
||||
except ImportError:
|
||||
# Define stubs if not available
|
||||
def create_memory_kv(*args, **kwargs):
|
||||
pass
|
||||
|
||||
def create_memory_vector(*args, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
try:
|
||||
from hanzo_agents import sequential_router, conditional_router, state_based_router
|
||||
except ImportError:
|
||||
try:
|
||||
# Try core module import
|
||||
from hanzo_agents.core.router import (
|
||||
sequential_router,
|
||||
conditional_router,
|
||||
state_based_router,
|
||||
)
|
||||
except ImportError:
|
||||
# Define stubs if not available
|
||||
def sequential_router(*args, **kwargs):
|
||||
pass
|
||||
|
||||
def conditional_router(*args, **kwargs):
|
||||
pass
|
||||
|
||||
def state_based_router(*args, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
try:
|
||||
from hanzo_agents.core.cli_agent import (
|
||||
GrokAgent,
|
||||
GeminiAgent,
|
||||
ClaudeCodeAgent,
|
||||
OpenAICodexAgent,
|
||||
)
|
||||
except ImportError:
|
||||
# Define stub classes if not available
|
||||
class ClaudeCodeAgent(Agent):
|
||||
pass
|
||||
|
||||
class OpenAICodexAgent(Agent):
|
||||
pass
|
||||
|
||||
class GeminiAgent(Agent):
|
||||
pass
|
||||
|
||||
class GrokAgent(Agent):
|
||||
pass
|
||||
|
||||
|
||||
from hanzo_mcp.tools.jupyter import get_read_only_jupyter_tools
|
||||
from hanzo_mcp.tools.filesystem import Edit, MultiEdit, get_read_only_filesystem_tools
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.agent.agent_tool import MCPAgent
|
||||
from hanzo_mcp.tools.common.batch_tool import BatchTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
|
||||
class AgentNode(TypedDict):
|
||||
"""Node in the agent network."""
|
||||
|
||||
id: str
|
||||
query: str
|
||||
model: Optional[str]
|
||||
role: Optional[str]
|
||||
connections: Optional[List[str]]
|
||||
receives_from: Optional[List[str]]
|
||||
file_path: Optional[str]
|
||||
|
||||
|
||||
class SwarmConfig(TypedDict):
|
||||
"""Configuration for an agent network."""
|
||||
|
||||
agents: Dict[str, AgentNode]
|
||||
entry_point: Optional[str]
|
||||
topology: Optional[str]
|
||||
|
||||
|
||||
class SwarmToolParams(TypedDict):
|
||||
"""Parameters for the SwarmTool."""
|
||||
|
||||
config: SwarmConfig
|
||||
query: str
|
||||
context: Optional[str]
|
||||
max_concurrent: Optional[int]
|
||||
use_memory: Optional[bool]
|
||||
memory_backend: Optional[str]
|
||||
|
||||
|
||||
class SwarmState(State):
|
||||
"""State for swarm execution."""
|
||||
|
||||
def __init__(
|
||||
self, config: SwarmConfig, initial_query: str, context: Optional[str] = None
|
||||
):
|
||||
"""Initialize swarm state."""
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.initial_query = initial_query
|
||||
self.context = context
|
||||
self.agent_results = {}
|
||||
self.completed_agents = set()
|
||||
self.current_agent = None
|
||||
self.execution_order = []
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
base_dict = super().to_dict()
|
||||
base_dict.update(
|
||||
{
|
||||
"config": self.config,
|
||||
"initial_query": self.initial_query,
|
||||
"context": self.context,
|
||||
"agent_results": self.agent_results,
|
||||
"completed_agents": list(self.completed_agents),
|
||||
"current_agent": self.current_agent,
|
||||
"execution_order": self.execution_order,
|
||||
}
|
||||
)
|
||||
return base_dict
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "SwarmState":
|
||||
"""Create from dictionary."""
|
||||
state = cls(
|
||||
config=data.get("config", {}),
|
||||
initial_query=data.get("initial_query", ""),
|
||||
context=data.get("context"),
|
||||
)
|
||||
state.agent_results = data.get("agent_results", {})
|
||||
state.completed_agents = set(data.get("completed_agents", []))
|
||||
state.current_agent = data.get("current_agent")
|
||||
state.execution_order = data.get("execution_order", [])
|
||||
return state
|
||||
|
||||
|
||||
class SwarmAgent(MCPAgent):
|
||||
"""Agent that executes within a swarm network."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent_id: str,
|
||||
agent_config: AgentNode,
|
||||
available_tools: List[BaseTool],
|
||||
permission_manager: PermissionManager,
|
||||
ctx: MCPContext,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize swarm agent."""
|
||||
# Set name and description from config
|
||||
self.name = agent_id
|
||||
self.description = agent_config.get("role", f"Agent {agent_id}")
|
||||
self.agent_config = agent_config
|
||||
|
||||
# Initialize with specified model
|
||||
model = agent_config.get("model")
|
||||
if model:
|
||||
model = self._normalize_model(model)
|
||||
else:
|
||||
model = "model://anthropic/claude-3-5-sonnet-20241022"
|
||||
|
||||
super().__init__(
|
||||
available_tools=available_tools,
|
||||
permission_manager=permission_manager,
|
||||
ctx=ctx,
|
||||
model=model,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def _normalize_model(self, model: str) -> str:
|
||||
"""Normalize model names to full format."""
|
||||
model_map = {
|
||||
"claude-3-5-sonnet": "model://anthropic/claude-3-5-sonnet-20241022",
|
||||
"claude-3-opus": "model://anthropic/claude-3-opus-20240229",
|
||||
"gpt-4o": "model://openai/gpt-4o",
|
||||
"gpt-4": "model://openai/gpt-4",
|
||||
"gemini-1.5-pro": "model://google/gemini-1.5-pro",
|
||||
"gemini-1.5-flash": "model://google/gemini-1.5-flash",
|
||||
}
|
||||
|
||||
# Check if it's already a model:// URI
|
||||
if model.startswith("model://"):
|
||||
return model
|
||||
|
||||
# Check mapping
|
||||
if model in model_map:
|
||||
return model_map[model]
|
||||
|
||||
# Assume it's a provider/model format
|
||||
if "/" in model:
|
||||
return f"model://{model}"
|
||||
|
||||
# Default to anthropic
|
||||
return f"model://anthropic/{model}"
|
||||
|
||||
async def run(
|
||||
self, state: SwarmState, history: History, network: Network
|
||||
) -> InferenceResult:
|
||||
"""Execute the swarm agent."""
|
||||
# Build prompt with context
|
||||
prompt_parts = []
|
||||
|
||||
# Add role context
|
||||
if self.agent_config.get("role"):
|
||||
prompt_parts.append(f"Your role: {self.agent_config['role']}")
|
||||
|
||||
# Add shared context
|
||||
if state.context:
|
||||
prompt_parts.append(f"Context:\n{state.context}")
|
||||
|
||||
# Add inputs from connected agents
|
||||
receives_from = self.agent_config.get("receives_from", [])
|
||||
if receives_from:
|
||||
inputs = {}
|
||||
for agent_id in receives_from:
|
||||
if agent_id in state.agent_results:
|
||||
inputs[agent_id] = state.agent_results[agent_id]
|
||||
|
||||
if inputs:
|
||||
prompt_parts.append("Input from previous agents:")
|
||||
for input_agent, input_result in inputs.items():
|
||||
prompt_parts.append(f"\n--- From {input_agent} ---\n{input_result}")
|
||||
|
||||
# Add file context if specified
|
||||
if self.agent_config.get("file_path"):
|
||||
prompt_parts.append(f"\nFile to work on: {self.agent_config['file_path']}")
|
||||
|
||||
# Add the main query
|
||||
prompt_parts.append(f"\nTask: {self.agent_config['query']}")
|
||||
|
||||
# Add initial query if this is entry point
|
||||
if state.current_agent == state.config.get("entry_point"):
|
||||
prompt_parts.append(f"\nMain objective: {state.initial_query}")
|
||||
|
||||
full_prompt = "\n\n".join(prompt_parts)
|
||||
|
||||
# Execute using base class
|
||||
messages = [
|
||||
{"role": "system", "content": self._get_system_prompt()},
|
||||
{"role": "user", "content": full_prompt},
|
||||
]
|
||||
|
||||
# Call model
|
||||
from hanzo_agents import ModelRegistry
|
||||
|
||||
adapter = ModelRegistry.get_adapter(self.model)
|
||||
response = await adapter.chat(messages)
|
||||
|
||||
# Store result in state
|
||||
state.agent_results[self.name] = response
|
||||
state.completed_agents.add(self.name)
|
||||
state.execution_order.append(self.name)
|
||||
|
||||
# Return result
|
||||
return InferenceResult(
|
||||
agent=self.name,
|
||||
content=response,
|
||||
metadata={
|
||||
"agent_id": self.name,
|
||||
"role": self.agent_config.get("role"),
|
||||
"connections": self.agent_config.get("connections", []),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class SwarmRouter(DeterministicRouter):
|
||||
"""Router for swarm agent orchestration."""
|
||||
|
||||
def __init__(self, swarm_config: SwarmConfig):
|
||||
"""Initialize swarm router."""
|
||||
self.swarm_config = swarm_config
|
||||
self.agents_config = swarm_config["agents"]
|
||||
self.entry_point = swarm_config.get("entry_point")
|
||||
|
||||
# Build dependency graph
|
||||
self.dependencies = {}
|
||||
self.dependents = {}
|
||||
|
||||
for agent_id, config in self.agents_config.items():
|
||||
# Dependencies (agents this one waits for)
|
||||
self.dependencies[agent_id] = set(config.get("receives_from", []))
|
||||
|
||||
# Dependents (agents that wait for this one)
|
||||
connections = config.get("connections", [])
|
||||
for conn in connections:
|
||||
if conn not in self.dependents:
|
||||
self.dependents[conn] = set()
|
||||
self.dependents[conn].add(agent_id)
|
||||
|
||||
def route(self, network, call_count, last_result, agent_stack):
|
||||
"""Determine next agent to execute."""
|
||||
state = network.state
|
||||
|
||||
# First call - start with entry point or roots
|
||||
if call_count == 0:
|
||||
if self.entry_point:
|
||||
state.current_agent = self.entry_point
|
||||
return self._get_agent_class(self.entry_point, agent_stack)
|
||||
else:
|
||||
# Find roots (no dependencies)
|
||||
roots = [aid for aid, deps in self.dependencies.items() if not deps]
|
||||
if roots:
|
||||
state.current_agent = roots[0]
|
||||
return self._get_agent_class(roots[0], agent_stack)
|
||||
|
||||
# Find next agent to execute
|
||||
for agent_id in self.agents_config:
|
||||
if agent_id in state.completed_agents:
|
||||
continue
|
||||
|
||||
# Check if all dependencies are met
|
||||
deps = self.dependencies.get(agent_id, set())
|
||||
if deps.issubset(state.completed_agents):
|
||||
state.current_agent = agent_id
|
||||
return self._get_agent_class(agent_id, agent_stack)
|
||||
|
||||
# No more agents to execute
|
||||
return None
|
||||
|
||||
def _get_agent_class(
|
||||
self, agent_id: str, agent_stack: List[type[Agent]]
|
||||
) -> type[Agent]:
|
||||
"""Get agent class for given agent ID."""
|
||||
# Find matching agent by name
|
||||
for agent_class in agent_stack:
|
||||
if hasattr(agent_class, "name") and agent_class.name == agent_id:
|
||||
return agent_class
|
||||
|
||||
# Not found - this shouldn't happen
|
||||
return None
|
||||
|
||||
|
||||
@final
|
||||
class SwarmTool(BaseTool):
|
||||
"""Tool for executing agent networks using hanzo-agents SDK."""
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "swarm"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Execute a network of AI agents with flexible connection topologies.
|
||||
|
||||
This tool enables sophisticated agent orchestration where agents can be connected
|
||||
in various network patterns. Each agent can pass results to connected agents,
|
||||
enabling complex workflows.
|
||||
|
||||
Features:
|
||||
- Flexible agent networks (tree, DAG, pipeline, star, mesh)
|
||||
- Each agent can use different models (Claude, GPT-4, Gemini, etc.)
|
||||
- Agents automatically pass results to connected agents
|
||||
- Parallel execution with dependency management
|
||||
- Full editing capabilities for each agent
|
||||
- Memory and state management via hanzo-agents SDK
|
||||
|
||||
Common Topologies:
|
||||
1. Tree (Architect pattern):
|
||||
architect → [frontend, backend, database] → reviewer
|
||||
|
||||
2. Pipeline (Sequential processing):
|
||||
analyzer → planner → implementer → tester → reviewer
|
||||
|
||||
3. Star (Central coordinator):
|
||||
coordinator ← → [agent1, agent2, agent3, agent4]
|
||||
|
||||
4. DAG (Complex dependencies):
|
||||
Multiple agents with custom connections
|
||||
|
||||
Models can be specified as:
|
||||
- Full: 'anthropic/claude-3-5-sonnet-20241022'
|
||||
- Short: 'claude-3-5-sonnet', 'gpt-4o', 'gemini-1.5-pro'
|
||||
- CLI tools: 'claude_cli', 'codex_cli', 'gemini_cli', 'grok_cli'
|
||||
- Model URIs: 'model://anthropic/claude-3-opus'
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
model: str | None = None,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
agent_max_iterations: int = 10,
|
||||
agent_max_tool_uses: int = 30,
|
||||
):
|
||||
"""Initialize the swarm tool."""
|
||||
self.permission_manager = permission_manager
|
||||
# Default to latest Claude Sonnet if no model specified
|
||||
from hanzo_mcp.tools.agent.code_auth import get_latest_claude_model
|
||||
|
||||
self.model = model or f"anthropic/{get_latest_claude_model()}"
|
||||
self.api_key = (
|
||||
api_key
|
||||
or os.environ.get("ANTHROPIC_API_KEY")
|
||||
or os.environ.get("CLAUDE_API_KEY")
|
||||
)
|
||||
self.base_url = base_url
|
||||
self.max_tokens = max_tokens
|
||||
self.agent_max_iterations = agent_max_iterations
|
||||
self.agent_max_tool_uses = agent_max_tool_uses
|
||||
|
||||
# Set up available tools for agents
|
||||
self.available_tools: list[BaseTool] = []
|
||||
self.available_tools.extend(
|
||||
get_read_only_filesystem_tools(self.permission_manager)
|
||||
)
|
||||
self.available_tools.extend(
|
||||
get_read_only_jupyter_tools(self.permission_manager)
|
||||
)
|
||||
|
||||
# Add edit tools
|
||||
self.available_tools.append(Edit(self.permission_manager))
|
||||
self.available_tools.append(MultiEdit(self.permission_manager))
|
||||
|
||||
# Add batch tool
|
||||
self.available_tools.append(
|
||||
BatchTool({t.name: t for t in self.available_tools})
|
||||
)
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[SwarmToolParams],
|
||||
) -> str:
|
||||
"""Execute the swarm tool."""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
config = params.get("config", {})
|
||||
initial_query = params.get("query", "")
|
||||
context = params.get("context", "")
|
||||
max_concurrent = params.get("max_concurrent", 10)
|
||||
use_memory = params.get("use_memory", False)
|
||||
memory_backend = params.get("memory_backend", "sqlite")
|
||||
|
||||
agents_config = config.get("agents", {})
|
||||
|
||||
if not agents_config:
|
||||
await tool_ctx.error("No agents provided")
|
||||
return "Error: At least one agent must be provided."
|
||||
|
||||
# hanzo-agents SDK is required (already imported above)
|
||||
|
||||
await tool_ctx.info(
|
||||
f"Starting swarm execution with {len(agents_config)} agents using hanzo-agents SDK"
|
||||
)
|
||||
|
||||
# Create state
|
||||
state = SwarmState(config=config, initial_query=initial_query, context=context)
|
||||
|
||||
# Create agent classes dynamically
|
||||
agent_classes = []
|
||||
for agent_id, agent_config in agents_config.items():
|
||||
# Check for CLI agents
|
||||
model = agent_config.get("model", self.model)
|
||||
|
||||
cli_agents = {
|
||||
"claude_cli": ClaudeCodeAgent,
|
||||
"codex_cli": OpenAICodexAgent,
|
||||
"gemini_cli": GeminiAgent,
|
||||
"grok_cli": GrokAgent,
|
||||
}
|
||||
|
||||
if model in cli_agents:
|
||||
# Use CLI agent
|
||||
agent_class = type(
|
||||
f"Swarm{agent_id}",
|
||||
(cli_agents[model],),
|
||||
{
|
||||
"name": agent_id,
|
||||
"description": agent_config.get("role", f"Agent {agent_id}"),
|
||||
"agent_config": agent_config,
|
||||
},
|
||||
)
|
||||
else:
|
||||
# Create dynamic SwarmAgent class
|
||||
agent_class = type(
|
||||
f"Swarm{agent_id}",
|
||||
(SwarmAgent,),
|
||||
{
|
||||
"name": agent_id,
|
||||
"__init__": lambda self, aid=agent_id, acfg=agent_config: SwarmAgent.__init__(
|
||||
self,
|
||||
agent_id=aid,
|
||||
agent_config=acfg,
|
||||
available_tools=self.available_tools,
|
||||
permission_manager=self.permission_manager,
|
||||
ctx=ctx,
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
agent_classes.append(agent_class)
|
||||
|
||||
# Create memory if requested
|
||||
memory_kv = None
|
||||
memory_vector = None
|
||||
if use_memory:
|
||||
memory_kv = create_memory_kv(memory_backend)
|
||||
memory_vector = create_memory_vector("simple")
|
||||
|
||||
# Create router
|
||||
router = SwarmRouter(config)
|
||||
|
||||
# Create network
|
||||
network = Network(
|
||||
state=state,
|
||||
agents=agent_classes,
|
||||
router=router,
|
||||
memory_kv=memory_kv,
|
||||
memory_vector=memory_vector,
|
||||
max_steps=self.agent_max_iterations * len(agents_config),
|
||||
)
|
||||
|
||||
# Execute
|
||||
try:
|
||||
final_state = await network.run()
|
||||
|
||||
# Format results
|
||||
return self._format_network_results(
|
||||
agents_config,
|
||||
final_state.agent_results,
|
||||
final_state.execution_order,
|
||||
config.get("entry_point"),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Swarm execution failed: {str(e)}")
|
||||
return f"Error: {str(e)}"
|
||||
|
||||
def _format_network_results(
|
||||
self,
|
||||
agents_config: Dict[str, Any],
|
||||
results: Dict[str, str],
|
||||
execution_order: List[str],
|
||||
entry_point: Optional[str],
|
||||
) -> str:
|
||||
"""Format results from agent network execution."""
|
||||
output = ["Agent Network Execution Results (hanzo-agents SDK)"]
|
||||
output.append("=" * 80)
|
||||
output.append(f"Total agents: {len(agents_config)}")
|
||||
output.append(f"Completed: {len(results)}")
|
||||
output.append(
|
||||
f"Failed: {len([r for r in results.values() if r.startswith('Error:')])}"
|
||||
)
|
||||
|
||||
if entry_point:
|
||||
output.append(f"Entry point: {entry_point}")
|
||||
|
||||
output.append(f"\nExecution Order: {' → '.join(execution_order)}")
|
||||
output.append("-" * 40)
|
||||
|
||||
# Detailed results
|
||||
output.append("\n\nDetailed Results:")
|
||||
output.append("=" * 80)
|
||||
|
||||
for agent_id in execution_order:
|
||||
if agent_id in results:
|
||||
config = agents_config.get(agent_id, {})
|
||||
role = config.get("role", "Agent")
|
||||
model = config.get("model", "default")
|
||||
|
||||
output.append(f"\n### {agent_id} ({role}) [{model}]")
|
||||
output.append("-" * 40)
|
||||
|
||||
result = results[agent_id]
|
||||
if result.startswith("Error:"):
|
||||
output.append(result)
|
||||
else:
|
||||
# Show first part of result
|
||||
lines = result.split("\n")
|
||||
preview_lines = lines[:10]
|
||||
output.extend(preview_lines)
|
||||
|
||||
if len(lines) > 10:
|
||||
output.append(f"... ({len(lines) - 10} more lines)")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this swarm tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def swarm(
|
||||
ctx: MCPContext,
|
||||
config: dict[str, Any],
|
||||
query: str,
|
||||
context: Optional[str] = None,
|
||||
max_concurrent: int = 10,
|
||||
use_memory: bool = False,
|
||||
memory_backend: str = "sqlite",
|
||||
) -> str:
|
||||
# Convert to typed format
|
||||
typed_config = SwarmConfig(
|
||||
agents=config.get("agents", {}),
|
||||
entry_point=config.get("entry_point"),
|
||||
topology=config.get("topology"),
|
||||
)
|
||||
|
||||
return await tool_self.call(
|
||||
ctx,
|
||||
config=typed_config,
|
||||
query=query,
|
||||
context=context,
|
||||
max_concurrent=max_concurrent,
|
||||
use_memory=use_memory,
|
||||
memory_backend=memory_backend,
|
||||
)
|
||||
@@ -1,577 +0,0 @@
|
||||
"""Swarm tool implementation for parallel and hierarchical agent execution.
|
||||
|
||||
This module implements the SwarmTool that enables both parallel execution of multiple
|
||||
agent instances and hierarchical workflows with specialized roles.
|
||||
"""
|
||||
|
||||
import os
|
||||
import asyncio
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
List,
|
||||
Unpack,
|
||||
Optional,
|
||||
TypedDict,
|
||||
final,
|
||||
override,
|
||||
)
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.agent.agent_tool import AgentTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
|
||||
class AgentNode(TypedDict):
|
||||
"""Node in the agent network.
|
||||
|
||||
Attributes:
|
||||
id: Unique identifier for this agent
|
||||
query: The specific query/task for this agent
|
||||
model: Optional model override (e.g., 'claude-3-5-sonnet', 'gpt-4o')
|
||||
role: Optional role description (e.g., 'architect', 'frontend', 'reviewer')
|
||||
connections: List of agent IDs this agent connects to (sends results to)
|
||||
receives_from: Optional list of agent IDs this agent receives input from
|
||||
file_path: Optional specific file for the agent to work on
|
||||
"""
|
||||
|
||||
id: str
|
||||
query: str
|
||||
model: Optional[str]
|
||||
role: Optional[str]
|
||||
connections: Optional[List[str]]
|
||||
receives_from: Optional[List[str]]
|
||||
file_path: Optional[str]
|
||||
|
||||
|
||||
class SwarmConfig(TypedDict):
|
||||
"""Configuration for an agent network.
|
||||
|
||||
Attributes:
|
||||
agents: Dictionary of agent configurations keyed by ID
|
||||
entry_point: ID of the first agent to execute (optional, defaults to finding roots)
|
||||
topology: Optional topology type (tree, dag, pipeline, star, mesh)
|
||||
"""
|
||||
|
||||
agents: Dict[str, AgentNode]
|
||||
entry_point: Optional[str]
|
||||
topology: Optional[str]
|
||||
|
||||
|
||||
class SwarmToolParams(TypedDict):
|
||||
"""Parameters for the SwarmTool.
|
||||
|
||||
Attributes:
|
||||
config: Agent network configuration
|
||||
query: Initial query to send to entry point agent(s)
|
||||
context: Optional context shared by all agents
|
||||
max_concurrent: Maximum number of concurrent agents (default: 10)
|
||||
"""
|
||||
|
||||
config: SwarmConfig
|
||||
query: str
|
||||
context: Optional[str]
|
||||
max_concurrent: Optional[int]
|
||||
|
||||
|
||||
@final
|
||||
class SwarmTool(BaseTool):
|
||||
"""Tool for executing multiple agent tasks in parallel.
|
||||
|
||||
The SwarmTool enables efficient parallel processing of multiple files or tasks
|
||||
by spawning independent agent instances for each task.
|
||||
"""
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "swarm"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Execute a network of AI agents with flexible connection topologies.
|
||||
|
||||
This tool enables sophisticated agent orchestration where agents can be connected
|
||||
in various network patterns. Each agent can pass results to connected agents,
|
||||
enabling complex workflows.
|
||||
|
||||
Features:
|
||||
- Flexible agent networks (tree, DAG, pipeline, star, mesh)
|
||||
- Each agent can use different models (Claude, GPT-4, Gemini, etc.)
|
||||
- Agents automatically pass results to connected agents
|
||||
- Parallel execution with dependency management
|
||||
- Full editing capabilities for each agent
|
||||
|
||||
Common Topologies:
|
||||
|
||||
1. Tree (Architect pattern):
|
||||
architect → [frontend, backend, database] → reviewer
|
||||
|
||||
2. Pipeline (Sequential processing):
|
||||
analyzer → planner → implementer → tester → reviewer
|
||||
|
||||
3. Star (Central coordinator):
|
||||
coordinator ← → [agent1, agent2, agent3, agent4]
|
||||
|
||||
4. DAG (Complex dependencies):
|
||||
Multiple agents with custom connections
|
||||
|
||||
Usage Example:
|
||||
|
||||
swarm(
|
||||
config={
|
||||
"agents": {
|
||||
"architect": {
|
||||
"id": "architect",
|
||||
"query": "Analyze codebase and create refactoring plan",
|
||||
"model": "claude-3-5-sonnet",
|
||||
"connections": ["frontend", "backend", "database"]
|
||||
},
|
||||
"frontend": {
|
||||
"id": "frontend",
|
||||
"query": "Refactor UI components based on architect's plan",
|
||||
"role": "Frontend Developer",
|
||||
"connections": ["reviewer"]
|
||||
},
|
||||
"backend": {
|
||||
"id": "backend",
|
||||
"query": "Refactor API endpoints based on architect's plan",
|
||||
"role": "Backend Developer",
|
||||
"connections": ["reviewer"]
|
||||
},
|
||||
"database": {
|
||||
"id": "database",
|
||||
"query": "Optimize database schema based on architect's plan",
|
||||
"role": "Database Expert",
|
||||
"connections": ["reviewer"]
|
||||
},
|
||||
"reviewer": {
|
||||
"id": "reviewer",
|
||||
"query": "Review all changes and ensure consistency",
|
||||
"model": "gpt-4o",
|
||||
"receives_from": ["frontend", "backend", "database"]
|
||||
}
|
||||
},
|
||||
"entry_point": "architect"
|
||||
},
|
||||
query="Refactor the authentication system for better security and performance"
|
||||
)
|
||||
|
||||
Models can be specified as:
|
||||
- Full: 'anthropic/claude-3-5-sonnet-20241022'
|
||||
- Short: 'claude-3-5-sonnet', 'gpt-4o', 'gemini-1.5-pro'
|
||||
- CLI tools: 'claude_cli', 'codex_cli', 'gemini_cli', 'grok_cli'
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
model: str | None = None,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
agent_max_iterations: int = 10,
|
||||
agent_max_tool_uses: int = 30,
|
||||
):
|
||||
"""Initialize the swarm tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
model: Optional model name override (defaults to Claude Sonnet)
|
||||
api_key: Optional API key for the model provider
|
||||
base_url: Optional base URL for the model provider
|
||||
max_tokens: Optional maximum tokens for model responses
|
||||
agent_max_iterations: Max iterations per agent (default: 10)
|
||||
agent_max_tool_uses: Max tool uses per agent (default: 30)
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
# Default to latest Claude Sonnet if no model specified
|
||||
from hanzo_mcp.tools.agent.code_auth import get_latest_claude_model
|
||||
|
||||
self.model = model or f"anthropic/{get_latest_claude_model()}"
|
||||
self.api_key = (
|
||||
api_key
|
||||
or os.environ.get("ANTHROPIC_API_KEY")
|
||||
or os.environ.get("CLAUDE_API_KEY")
|
||||
)
|
||||
self.base_url = base_url
|
||||
self.max_tokens = max_tokens
|
||||
self.agent_max_iterations = agent_max_iterations
|
||||
self.agent_max_tool_uses = agent_max_tool_uses
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[SwarmToolParams],
|
||||
) -> str:
|
||||
"""Execute the swarm tool.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Combined results from all agents
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
agents = params.get("agents", [])
|
||||
manager_query = params.get("manager_query")
|
||||
reviewer_query = params.get("reviewer_query")
|
||||
common_context = params.get("common_context", "")
|
||||
max_concurrent = params.get("max_concurrent", 10)
|
||||
|
||||
if not agents:
|
||||
await tool_ctx.error("No agents provided")
|
||||
return "Error: At least one agent must be provided."
|
||||
|
||||
# Extract parameters
|
||||
config = params.get("config", {})
|
||||
initial_query = params.get("query", "")
|
||||
context = params.get("context", "")
|
||||
|
||||
agents_config = config.get("agents", {})
|
||||
entry_point = config.get("entry_point")
|
||||
|
||||
await tool_ctx.info(
|
||||
f"Starting swarm execution with {len(agents_config)} agents"
|
||||
)
|
||||
|
||||
# Build agent network
|
||||
agent_instances = {}
|
||||
agent_results = {}
|
||||
execution_queue = asyncio.Queue()
|
||||
completed_agents = set()
|
||||
|
||||
# Create agent instances
|
||||
for agent_id, agent_config in agents_config.items():
|
||||
model = agent_config.get("model", self.model)
|
||||
|
||||
# Support CLI tools
|
||||
cli_tools = {
|
||||
"claude_cli": self._get_cli_tool("claude_cli"),
|
||||
"codex_cli": self._get_cli_tool("codex_cli"),
|
||||
"gemini_cli": self._get_cli_tool("gemini_cli"),
|
||||
"grok_cli": self._get_cli_tool("grok_cli"),
|
||||
}
|
||||
|
||||
if model in cli_tools:
|
||||
agent = cli_tools[model]
|
||||
else:
|
||||
# Regular agent with model
|
||||
agent = AgentTool(
|
||||
permission_manager=self.permission_manager,
|
||||
model=self._normalize_model(model),
|
||||
api_key=self.api_key,
|
||||
base_url=self.base_url,
|
||||
max_tokens=self.max_tokens,
|
||||
max_iterations=self.agent_max_iterations,
|
||||
max_tool_uses=self.agent_max_tool_uses,
|
||||
)
|
||||
|
||||
agent_instances[agent_id] = agent
|
||||
|
||||
# Find entry points (agents with no incoming connections)
|
||||
if entry_point:
|
||||
await execution_queue.put((entry_point, initial_query, {}))
|
||||
else:
|
||||
# Find root agents (no receives_from)
|
||||
roots = []
|
||||
for agent_id, agent_config in agents_config.items():
|
||||
if not agent_config.get("receives_from"):
|
||||
# Check if any other agent connects to this one
|
||||
has_incoming = False
|
||||
for other_config in agents_config.values():
|
||||
if (
|
||||
other_config.get("connections")
|
||||
and agent_id in other_config["connections"]
|
||||
):
|
||||
has_incoming = True
|
||||
break
|
||||
if not has_incoming:
|
||||
roots.append(agent_id)
|
||||
|
||||
if not roots:
|
||||
await tool_ctx.error("No entry point found in agent network")
|
||||
return "Error: Could not determine entry point for agent network"
|
||||
|
||||
for root in roots:
|
||||
await execution_queue.put((root, initial_query, {}))
|
||||
|
||||
# Execute agents in network order
|
||||
async def execute_agent(
|
||||
agent_id: str, query: str, inputs: Dict[str, str]
|
||||
) -> str:
|
||||
"""Execute a single agent in the network."""
|
||||
async with semaphore:
|
||||
try:
|
||||
agent_config = agents_config[agent_id]
|
||||
agent = agent_instances[agent_id]
|
||||
|
||||
await tool_ctx.info(
|
||||
f"Executing agent: {agent_id} ({agent_config.get('role', 'Agent')})"
|
||||
)
|
||||
|
||||
# Build prompt with context and inputs
|
||||
prompt_parts = []
|
||||
|
||||
# Add role context
|
||||
if agent_config.get("role"):
|
||||
prompt_parts.append(f"Your role: {agent_config['role']}")
|
||||
|
||||
# Add shared context
|
||||
if context:
|
||||
prompt_parts.append(f"Context:\n{context}")
|
||||
|
||||
# Add inputs from connected agents
|
||||
if inputs:
|
||||
prompt_parts.append("Input from previous agents:")
|
||||
for input_agent, input_result in inputs.items():
|
||||
prompt_parts.append(
|
||||
f"\n--- From {input_agent} ---\n{input_result}"
|
||||
)
|
||||
|
||||
# Add file context if specified
|
||||
if agent_config.get("file_path"):
|
||||
prompt_parts.append(
|
||||
f"\nFile to work on: {agent_config['file_path']}"
|
||||
)
|
||||
|
||||
# Add the main query
|
||||
prompt_parts.append(f"\nTask: {agent_config['query']}")
|
||||
|
||||
# Combine query with initial query if this is entry point
|
||||
if query and query != agent_config["query"]:
|
||||
prompt_parts.append(f"\nMain objective: {query}")
|
||||
|
||||
full_prompt = "\n\n".join(prompt_parts)
|
||||
|
||||
# Execute the agent
|
||||
result = await agent.call(ctx, prompts=full_prompt)
|
||||
|
||||
await tool_ctx.info(f"Agent {agent_id} completed")
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Agent {agent_id} failed: {str(e)}"
|
||||
await tool_ctx.error(error_msg)
|
||||
return f"Error: {error_msg}"
|
||||
|
||||
# Process agent network
|
||||
running_tasks = set()
|
||||
|
||||
while not execution_queue.empty() or running_tasks:
|
||||
# Start new tasks up to concurrency limit
|
||||
while not execution_queue.empty() and len(running_tasks) < max_concurrent:
|
||||
agent_id, query, inputs = await execution_queue.get()
|
||||
|
||||
if agent_id not in completed_agents:
|
||||
# Check if all dependencies are met
|
||||
agent_config = agents_config[agent_id]
|
||||
receives_from = agent_config.get("receives_from", [])
|
||||
|
||||
# Collect inputs from dependencies
|
||||
ready = True
|
||||
for dep in receives_from:
|
||||
if dep not in agent_results:
|
||||
ready = False
|
||||
# Re-queue for later
|
||||
await execution_queue.put((agent_id, query, inputs))
|
||||
break
|
||||
else:
|
||||
inputs[dep] = agent_results[dep]
|
||||
|
||||
if ready:
|
||||
# Execute agent
|
||||
task = asyncio.create_task(
|
||||
execute_agent(agent_id, query, inputs)
|
||||
)
|
||||
running_tasks.add(task)
|
||||
|
||||
async def handle_completion(task, agent_id=agent_id):
|
||||
result = await task
|
||||
agent_results[agent_id] = result
|
||||
completed_agents.add(agent_id)
|
||||
running_tasks.discard(task)
|
||||
|
||||
# Queue connected agents
|
||||
agent_config = agents_config[agent_id]
|
||||
connections = agent_config.get("connections", [])
|
||||
for next_agent in connections:
|
||||
if next_agent in agents_config:
|
||||
await execution_queue.put(
|
||||
(next_agent, "", {agent_id: result})
|
||||
)
|
||||
|
||||
asyncio.create_task(handle_completion(task))
|
||||
|
||||
# Wait a bit if we're at capacity
|
||||
if running_tasks:
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Wait for all tasks to complete
|
||||
if running_tasks:
|
||||
await asyncio.gather(*running_tasks, return_exceptions=True)
|
||||
|
||||
# Format results
|
||||
return self._format_network_results(agents_config, agent_results, entry_point)
|
||||
|
||||
def _normalize_model(self, model: str) -> str:
|
||||
"""Normalize model names to full format."""
|
||||
model_map = {
|
||||
"claude-3-5-sonnet": "anthropic/claude-3-5-sonnet-20241022",
|
||||
"claude-3-opus": "anthropic/claude-3-opus-20240229",
|
||||
"gpt-4o": "openai/gpt-4o",
|
||||
"gpt-4": "openai/gpt-4",
|
||||
"gemini-1.5-pro": "google/gemini-1.5-pro",
|
||||
"gemini-1.5-flash": "google/gemini-1.5-flash",
|
||||
}
|
||||
return model_map.get(model, model)
|
||||
|
||||
def _get_cli_tool(self, tool_name: str):
|
||||
"""Get CLI tool instance."""
|
||||
# Import here to avoid circular imports
|
||||
if tool_name == "claude_cli":
|
||||
from hanzo_mcp.tools.agent.claude_cli_tool import ClaudeCLITool
|
||||
|
||||
return ClaudeCLITool(self.permission_manager)
|
||||
elif tool_name == "codex_cli":
|
||||
from hanzo_mcp.tools.agent.codex_cli_tool import CodexCLITool
|
||||
|
||||
return CodexCLITool(self.permission_manager)
|
||||
elif tool_name == "gemini_cli":
|
||||
from hanzo_mcp.tools.agent.gemini_cli_tool import GeminiCLITool
|
||||
|
||||
return GeminiCLITool(self.permission_manager)
|
||||
elif tool_name == "grok_cli":
|
||||
from hanzo_mcp.tools.agent.grok_cli_tool import GrokCLITool
|
||||
|
||||
return GrokCLITool(self.permission_manager)
|
||||
return None
|
||||
|
||||
def _format_network_results(
|
||||
self,
|
||||
agents_config: Dict[str, Any],
|
||||
results: Dict[str, str],
|
||||
entry_point: Optional[str],
|
||||
) -> str:
|
||||
"""Format results from agent network execution."""
|
||||
output = ["Agent Network Execution Results"]
|
||||
output.append("=" * 80)
|
||||
output.append(f"Total agents: {len(agents_config)}")
|
||||
output.append(f"Completed: {len(results)}")
|
||||
output.append(
|
||||
f"Failed: {len([r for r in results.values() if r.startswith('Error:')])}"
|
||||
)
|
||||
|
||||
if entry_point:
|
||||
output.append(f"Entry point: {entry_point}")
|
||||
|
||||
output.append("\nExecution Flow:")
|
||||
output.append("-" * 40)
|
||||
|
||||
# Show results in execution order
|
||||
def format_agent_tree(agent_id: str, level: int = 0) -> List[str]:
|
||||
lines = []
|
||||
indent = " " * level
|
||||
|
||||
if agent_id in agents_config:
|
||||
config = agents_config[agent_id]
|
||||
role = config.get("role", "Agent")
|
||||
model = config.get("model", "default")
|
||||
|
||||
status = (
|
||||
"✅"
|
||||
if agent_id in results
|
||||
and not results[agent_id].startswith("Error:")
|
||||
else "❌"
|
||||
)
|
||||
lines.append(f"{indent}{status} {agent_id} ({role}) [{model}]")
|
||||
|
||||
# Show connections
|
||||
connections = config.get("connections", [])
|
||||
for conn in connections:
|
||||
if conn in agents_config:
|
||||
lines.extend(format_agent_tree(conn, level + 1))
|
||||
|
||||
return lines
|
||||
|
||||
# Start from entry point or roots
|
||||
if entry_point:
|
||||
output.extend(format_agent_tree(entry_point))
|
||||
else:
|
||||
# Find roots
|
||||
roots = []
|
||||
for agent_id in agents_config:
|
||||
has_incoming = False
|
||||
for config in agents_config.values():
|
||||
if config.get("connections") and agent_id in config["connections"]:
|
||||
has_incoming = True
|
||||
break
|
||||
if not has_incoming:
|
||||
roots.append(agent_id)
|
||||
|
||||
for root in roots:
|
||||
output.extend(format_agent_tree(root))
|
||||
|
||||
# Detailed results
|
||||
output.append("\n\nDetailed Results:")
|
||||
output.append("=" * 80)
|
||||
|
||||
for agent_id, result in results.items():
|
||||
config = agents_config.get(agent_id, {})
|
||||
role = config.get("role", "Agent")
|
||||
|
||||
output.append(f"\n### {agent_id} ({role})")
|
||||
output.append("-" * 40)
|
||||
|
||||
if result.startswith("Error:"):
|
||||
output.append(result)
|
||||
else:
|
||||
# Show first part of result
|
||||
lines = result.split("\n")
|
||||
preview_lines = lines[:10]
|
||||
output.extend(preview_lines)
|
||||
|
||||
if len(lines) > 10:
|
||||
output.append(f"... ({len(lines) - 10} more lines)")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this swarm tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def swarm(
|
||||
ctx: MCPContext,
|
||||
config: dict[str, Any],
|
||||
query: str,
|
||||
context: Optional[str] = None,
|
||||
max_concurrent: int = 10,
|
||||
) -> str:
|
||||
# Convert to typed format
|
||||
typed_config = SwarmConfig(
|
||||
agents=config.get("agents", {}),
|
||||
entry_point=config.get("entry_point"),
|
||||
topology=config.get("topology"),
|
||||
)
|
||||
|
||||
return await tool_self.call(
|
||||
ctx,
|
||||
config=typed_config,
|
||||
query=query,
|
||||
context=context,
|
||||
max_concurrent=max_concurrent,
|
||||
)
|
||||
@@ -1,109 +0,0 @@
|
||||
"""Tool adapters for converting between MCP tools and OpenAI tools.
|
||||
|
||||
This module handles conversion between MCP tool formats and OpenAI function
|
||||
formats, making MCP tools available to the OpenAI API, and processing tool inputs
|
||||
and outputs for agent execution.
|
||||
"""
|
||||
|
||||
# Import litellm with warnings suppressed
|
||||
import warnings
|
||||
|
||||
from openai.types import FunctionParameters
|
||||
from openai.types.chat import ChatCompletionToolParam
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", DeprecationWarning)
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
|
||||
|
||||
def convert_tools_to_openai_functions(
|
||||
tools: list[BaseTool],
|
||||
) -> list[ChatCompletionToolParam]:
|
||||
"""Convert MCP tools to OpenAI function format.
|
||||
|
||||
Args:
|
||||
tools: List of MCP tools
|
||||
|
||||
Returns:
|
||||
List of tools formatted for OpenAI API
|
||||
"""
|
||||
openai_tools: list[ChatCompletionToolParam] = []
|
||||
for tool in tools:
|
||||
openai_tool: ChatCompletionToolParam = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool.name,
|
||||
"description": tool.description,
|
||||
"parameters": convert_tool_parameters(tool),
|
||||
},
|
||||
}
|
||||
openai_tools.append(openai_tool)
|
||||
return openai_tools
|
||||
|
||||
|
||||
def convert_tool_parameters(tool: BaseTool) -> FunctionParameters:
|
||||
"""Convert tool parameters to OpenAI format.
|
||||
|
||||
Args:
|
||||
tool: MCP tool
|
||||
|
||||
Returns:
|
||||
Parameter schema in OpenAI format
|
||||
"""
|
||||
# Start with a copy of the parameters
|
||||
params = tool.parameters.copy()
|
||||
|
||||
# Ensure the schema has the right format for OpenAI
|
||||
if "properties" not in params:
|
||||
params["properties"] = {}
|
||||
|
||||
if "type" not in params:
|
||||
params["type"] = "object"
|
||||
|
||||
if "required" not in params:
|
||||
params["required"] = tool.required
|
||||
|
||||
return params
|
||||
|
||||
|
||||
def supports_parallel_function_calling(model: str) -> bool:
|
||||
"""Check if a model supports parallel function calling.
|
||||
|
||||
Args:
|
||||
model: Model identifier in LiteLLM format (e.g., "openai/gpt-4-turbo-preview")
|
||||
|
||||
Returns:
|
||||
True if the model supports parallel function calling, False otherwise
|
||||
"""
|
||||
# Since litellm doesn't have this function, we'll implement a simple check
|
||||
# based on known models that support parallel function calling
|
||||
parallel_capable_models = {
|
||||
# OpenAI models that support parallel function calling
|
||||
"gpt-4-turbo",
|
||||
"gpt-4-turbo-preview",
|
||||
"gpt-4-turbo-2024-04-09",
|
||||
"gpt-4o",
|
||||
"gpt-4o-mini",
|
||||
"gpt-4o-2024-05-13",
|
||||
"gpt-4o-2024-08-06",
|
||||
"gpt-3.5-turbo",
|
||||
"gpt-3.5-turbo-0125",
|
||||
"gpt-3.5-turbo-1106",
|
||||
# Anthropic models with tool support
|
||||
"claude-3-opus",
|
||||
"claude-3-sonnet",
|
||||
"claude-3-haiku",
|
||||
"claude-3-5-sonnet",
|
||||
"claude-3-5-sonnet-20241022",
|
||||
}
|
||||
|
||||
# Extract model name without provider prefix
|
||||
model_name = model.split("/")[-1] if "/" in model else model
|
||||
|
||||
# Check if the base model name matches any known parallel-capable models
|
||||
for capable_model in parallel_capable_models:
|
||||
if model_name.startswith(capable_model):
|
||||
return True
|
||||
|
||||
return False
|
||||
@@ -1,259 +0,0 @@
|
||||
"""Unified CLI Tools - DRY implementation using base agent classes.
|
||||
|
||||
This module provides the single, clean implementation of all CLI tools
|
||||
following Python best practices and eliminating all duplication.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
from pathlib import Path
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context
|
||||
|
||||
from ..common.base import BaseTool
|
||||
from ...core.base_agent import CLIAgent, AgentConfig
|
||||
from ..common.permissions import PermissionManager
|
||||
from ...core.model_registry import registry
|
||||
|
||||
|
||||
class UnifiedCLITool(BaseTool, CLIAgent):
|
||||
"""Unified CLI tool that combines BaseTool and CLIAgent functionality.
|
||||
|
||||
MRO: BaseTool first for proper method resolution order.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
description: str,
|
||||
cli_command: str,
|
||||
default_model: str,
|
||||
permission_manager: Optional[PermissionManager] = None,
|
||||
):
|
||||
"""Initialize unified CLI tool.
|
||||
|
||||
Args:
|
||||
name: Tool name
|
||||
description: Tool description
|
||||
cli_command: CLI command to execute
|
||||
default_model: Default model to use
|
||||
permission_manager: Permission manager for access control
|
||||
"""
|
||||
# Initialize CLIAgent with config
|
||||
config = AgentConfig(model=default_model)
|
||||
CLIAgent.__init__(self, config)
|
||||
|
||||
# Store tool metadata
|
||||
self._name = name
|
||||
self._description = description
|
||||
self._cli_command = cli_command
|
||||
self.permission_manager = permission_manager
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return self._description
|
||||
|
||||
@property
|
||||
def cli_command(self) -> str:
|
||||
return self._cli_command
|
||||
|
||||
def build_command(self, prompt: str, **kwargs: Any) -> List[str]:
|
||||
"""Build the CLI command with model-specific formatting.
|
||||
|
||||
Args:
|
||||
prompt: The prompt
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
Command arguments list
|
||||
"""
|
||||
command = [self.cli_command]
|
||||
|
||||
# Get model config from registry
|
||||
model_config = registry.get(self.config.model)
|
||||
|
||||
# Handle different CLI tool formats
|
||||
if self.cli_command == "claude":
|
||||
if model_config:
|
||||
command.extend(["--model", model_config.full_name])
|
||||
# Claude takes prompt via stdin
|
||||
return command
|
||||
|
||||
elif self.cli_command == "openai":
|
||||
# OpenAI CLI format
|
||||
command.extend(["api", "chat.completions.create"])
|
||||
if model_config:
|
||||
command.extend(["-m", model_config.full_name])
|
||||
command.extend(["-g", "user", prompt])
|
||||
return command
|
||||
|
||||
elif self.cli_command in ["gemini", "grok"]:
|
||||
# Simple format: command --model MODEL prompt
|
||||
if model_config:
|
||||
command.extend(["--model", model_config.full_name])
|
||||
command.append(prompt)
|
||||
return command
|
||||
|
||||
elif self.cli_command == "openhands":
|
||||
# OpenHands format
|
||||
command.extend(["run", prompt])
|
||||
if model_config:
|
||||
command.extend(["--model", model_config.full_name])
|
||||
if self.config.working_dir:
|
||||
command.extend(["--workspace", str(self.config.working_dir)])
|
||||
return command
|
||||
|
||||
elif self.cli_command == "hanzo":
|
||||
# Hanzo dev format
|
||||
command.append("dev")
|
||||
if model_config:
|
||||
command.extend(["--model", model_config.full_name])
|
||||
command.extend(["--prompt", prompt])
|
||||
return command
|
||||
|
||||
elif self.cli_command == "cline":
|
||||
# Cline format
|
||||
command.append(prompt)
|
||||
command.append("--no-interactive")
|
||||
return command
|
||||
|
||||
elif self.cli_command == "aider":
|
||||
# Aider format
|
||||
if model_config:
|
||||
command.extend(["--model", model_config.full_name])
|
||||
command.extend(["--message", prompt])
|
||||
command.extend(["--yes", "--no-stream"])
|
||||
return command
|
||||
|
||||
elif self.cli_command == "ollama":
|
||||
# Ollama format for local models
|
||||
command.extend(["run", self.config.model.replace("ollama/", "")])
|
||||
command.append(prompt)
|
||||
return command
|
||||
|
||||
# Default format
|
||||
command.append(prompt)
|
||||
return command
|
||||
|
||||
async def call(self, ctx: Context[Any, Any, Any], **params: Any) -> str:
|
||||
"""Execute the CLI tool via MCP interface.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Execution result
|
||||
"""
|
||||
# Update config from params
|
||||
if params.get("model"):
|
||||
self.config.model = registry.resolve(params["model"])
|
||||
if params.get("working_dir"):
|
||||
self.config.working_dir = Path(params["working_dir"])
|
||||
if params.get("timeout"):
|
||||
self.config.timeout = params["timeout"]
|
||||
|
||||
# Execute using base agent
|
||||
result = await self.execute(
|
||||
params.get("prompt", ""),
|
||||
context=ctx,
|
||||
)
|
||||
|
||||
return result.content
|
||||
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this tool with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
"""
|
||||
tool_self = self
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def tool_wrapper(
|
||||
prompt: str,
|
||||
ctx: Context[Any, Any, Any],
|
||||
model: Optional[str] = None,
|
||||
working_dir: Optional[str] = None,
|
||||
timeout: int = 300,
|
||||
) -> str:
|
||||
return await tool_self.call(
|
||||
ctx,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
working_dir=working_dir,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
def create_cli_tools(permission_manager: Optional[PermissionManager] = None) -> Dict[str, UnifiedCLITool]:
|
||||
"""Create all CLI tools with unified implementation.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
|
||||
Returns:
|
||||
Dictionary of tool name to tool instance
|
||||
"""
|
||||
tools = {}
|
||||
|
||||
# Define all tools with their configurations
|
||||
tool_configs = [
|
||||
("claude", "Execute Claude CLI for AI assistance", "claude", "claude"),
|
||||
("cc", "Claude Code CLI (alias for claude)", "claude", "claude"),
|
||||
("codex", "Execute OpenAI Codex/GPT-4 CLI", "openai", "gpt-4-turbo"),
|
||||
("gemini", "Execute Google Gemini CLI", "gemini", "gemini"),
|
||||
("grok", "Execute xAI Grok CLI", "grok", "grok"),
|
||||
("openhands", "Execute OpenHands for autonomous coding", "openhands", "claude"),
|
||||
("oh", "OpenHands CLI (alias)", "openhands", "claude"),
|
||||
("hanzo_dev", "Execute Hanzo Dev AI assistant", "hanzo", "claude"),
|
||||
("cline", "Execute Cline for autonomous coding", "cline", "claude"),
|
||||
("aider", "Execute Aider for AI pair programming", "aider", "gpt-4-turbo"),
|
||||
]
|
||||
|
||||
for name, description, cli_command, default_model in tool_configs:
|
||||
tools[name] = UnifiedCLITool(
|
||||
name=name,
|
||||
description=description,
|
||||
cli_command=cli_command,
|
||||
default_model=default_model,
|
||||
permission_manager=permission_manager,
|
||||
)
|
||||
|
||||
return tools
|
||||
|
||||
|
||||
def register_cli_tools(
|
||||
mcp_server: FastMCP,
|
||||
permission_manager: Optional[PermissionManager] = None,
|
||||
) -> List[BaseTool]:
|
||||
"""Register all CLI tools with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
permission_manager: Permission manager for access control
|
||||
|
||||
Returns:
|
||||
List of registered CLI tools
|
||||
"""
|
||||
tools = create_cli_tools(permission_manager)
|
||||
|
||||
# Register each tool
|
||||
for tool in tools.values():
|
||||
tool.register(mcp_server)
|
||||
|
||||
return list(tools.values())
|
||||
|
||||
|
||||
__all__ = [
|
||||
"UnifiedCLITool",
|
||||
"create_cli_tools",
|
||||
"register_cli_tools",
|
||||
]
|
||||
@@ -1,45 +0,0 @@
|
||||
"""Common utilities for Hanzo AI tools."""
|
||||
|
||||
from mcp.server import FastMCP
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool, ToolRegistry
|
||||
from hanzo_mcp.tools.common.batch_tool import BatchTool
|
||||
from hanzo_mcp.tools.common.critic_tool import CriticTool
|
||||
from hanzo_mcp.tools.common.thinking_tool import ThinkingTool
|
||||
|
||||
|
||||
def register_thinking_tool(
|
||||
mcp_server: FastMCP,
|
||||
) -> list[BaseTool]:
|
||||
"""Register thinking tools with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
"""
|
||||
thinking_tool = ThinkingTool()
|
||||
ToolRegistry.register_tool(mcp_server, thinking_tool)
|
||||
return [thinking_tool]
|
||||
|
||||
|
||||
def register_critic_tool(
|
||||
mcp_server: FastMCP,
|
||||
) -> list[BaseTool]:
|
||||
"""Register critic tool with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
"""
|
||||
critic_tool = CriticTool()
|
||||
ToolRegistry.register_tool(mcp_server, critic_tool)
|
||||
return [critic_tool]
|
||||
|
||||
|
||||
def register_batch_tool(mcp_server: FastMCP, tools: dict[str, BaseTool]) -> None:
|
||||
"""Register batch tool with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
tools: Dictionary mapping tool names to tool instances
|
||||
"""
|
||||
batch_tool = BatchTool(tools)
|
||||
ToolRegistry.register_tool(mcp_server, batch_tool)
|
||||
@@ -1,189 +0,0 @@
|
||||
"""Base classes for Hanzo AI tools.
|
||||
|
||||
This module provides abstract base classes that define interfaces and common functionality
|
||||
for all tools used in Hanzo AI. These abstractions help ensure consistent tool
|
||||
behavior and provide a foundation for tool registration and management.
|
||||
"""
|
||||
|
||||
import functools
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Callable, final
|
||||
from collections.abc import Awaitable
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.validation import (
|
||||
ValidationResult,
|
||||
validate_path_parameter,
|
||||
)
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
|
||||
def handle_connection_errors(
|
||||
func: Callable[..., Awaitable[str]],
|
||||
) -> Callable[..., Awaitable[str]]:
|
||||
"""Decorator to handle connection errors in MCP tool functions.
|
||||
|
||||
This decorator wraps tool functions to catch ClosedResourceError and other
|
||||
connection-related exceptions that occur when the client disconnects.
|
||||
|
||||
Args:
|
||||
func: The async tool function to wrap
|
||||
|
||||
Returns:
|
||||
Wrapped function that handles connection errors gracefully
|
||||
"""
|
||||
|
||||
@functools.wraps(func)
|
||||
async def wrapper(*args: Any, **kwargs: Any) -> str:
|
||||
try:
|
||||
return await func(*args, **kwargs)
|
||||
except Exception as e:
|
||||
# Check if this is a connection-related error
|
||||
error_name = type(e).__name__
|
||||
if any(
|
||||
name in error_name
|
||||
for name in [
|
||||
"ClosedResourceError",
|
||||
"ConnectionError",
|
||||
"BrokenPipeError",
|
||||
]
|
||||
):
|
||||
# Client has disconnected - log the error but don't crash
|
||||
# Return a simple error message (though it likely won't be received)
|
||||
return f"Client disconnected during operation: {error_name}"
|
||||
else:
|
||||
# Re-raise non-connection errors
|
||||
raise
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
class BaseTool(ABC):
|
||||
"""Abstract base class for all Hanzo AI tools.
|
||||
|
||||
This class defines the core interface that all tools must implement, ensuring
|
||||
consistency in how tools are registered, documented, and called.
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
"""Get the tool name.
|
||||
|
||||
Returns:
|
||||
The tool name as it will appear in the MCP server
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def description(self) -> str:
|
||||
"""Get the tool description.
|
||||
|
||||
Returns:
|
||||
Detailed description of the tool's purpose and usage
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def call(self, ctx: MCPContext, **params: Any) -> Any:
|
||||
"""Execute the tool with the given parameters.
|
||||
|
||||
Args:
|
||||
ctx: MCP context for the tool call
|
||||
**params: Tool parameters provided by the caller
|
||||
|
||||
Returns:
|
||||
Tool execution result as a string
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this tool with the MCP server.
|
||||
|
||||
This method must be implemented by each tool class to create a wrapper function
|
||||
with explicitly defined parameters that calls this tool's call method.
|
||||
The wrapper function is then registered with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class FileSystemTool(BaseTool, ABC):
|
||||
"""Base class for filesystem-related tools.
|
||||
|
||||
Provides common functionality for working with files and directories,
|
||||
including permission checking and path validation.
|
||||
"""
|
||||
|
||||
def __init__(self, permission_manager: PermissionManager) -> None:
|
||||
"""Initialize filesystem tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
"""
|
||||
self.permission_manager: PermissionManager = permission_manager
|
||||
|
||||
def validate_path(self, path: str, param_name: str = "path") -> ValidationResult:
|
||||
"""Validate a path parameter.
|
||||
|
||||
Args:
|
||||
path: Path to validate
|
||||
param_name: Name of the parameter (for error messages)
|
||||
|
||||
Returns:
|
||||
Validation result containing validation status and error message if any
|
||||
"""
|
||||
return validate_path_parameter(path, param_name)
|
||||
|
||||
def is_path_allowed(self, path: str) -> bool:
|
||||
"""Check if a path is allowed according to permission settings.
|
||||
|
||||
Args:
|
||||
path: Path to check
|
||||
|
||||
Returns:
|
||||
True if the path is allowed, False otherwise
|
||||
"""
|
||||
return self.permission_manager.is_path_allowed(path)
|
||||
|
||||
|
||||
@final
|
||||
class ToolRegistry:
|
||||
"""Registry for Hanzo AI tools.
|
||||
|
||||
Provides functionality for registering tool implementations with an MCP server,
|
||||
handling the conversion between tool classes and MCP tool functions.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def register_tool(mcp_server: FastMCP, tool: BaseTool) -> None:
|
||||
"""Register a tool with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
tool: The tool to register
|
||||
"""
|
||||
# Check if tool is enabled before registering
|
||||
# Import here to avoid circular imports
|
||||
from hanzo_mcp.tools.common.tool_enable import ToolEnableTool
|
||||
|
||||
if ToolEnableTool.is_tool_enabled(tool.name):
|
||||
# Use the tool's register method which handles all the details
|
||||
tool.register(mcp_server)
|
||||
|
||||
@staticmethod
|
||||
def register_tools(mcp_server: FastMCP, tools: list[BaseTool]) -> None:
|
||||
"""Register multiple tools with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
tools: List of tools to register
|
||||
"""
|
||||
for tool in tools:
|
||||
ToolRegistry.register_tool(mcp_server, tool)
|
||||
@@ -1,429 +0,0 @@
|
||||
"""Batch tool implementation for Hanzo AI.
|
||||
|
||||
This module provides the BatchTool that allows for executing multiple tools in
|
||||
parallel or serial depending on their characteristics.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import Any, Unpack, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.truncate import truncate_response
|
||||
from hanzo_mcp.tools.common.fastmcp_pagination import (
|
||||
CursorData,
|
||||
create_paginated_response,
|
||||
)
|
||||
|
||||
|
||||
class InvocationItem(TypedDict):
|
||||
"""A single tool invocation item.
|
||||
|
||||
Attributes:
|
||||
tool_name: The name of the tool to invoke
|
||||
input: The input to pass to the tool
|
||||
"""
|
||||
|
||||
tool_name: Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="The name of the tool to invoke",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
input: Annotated[
|
||||
dict[str, Any],
|
||||
Field(
|
||||
description="The input to pass to the tool",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
Description = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="A short (3-5 word) description of the batch operation",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
Invocations = Annotated[
|
||||
list[InvocationItem],
|
||||
Field(
|
||||
description="The list of tool invocations to execute (required -- you MUST provide at least one tool invocation)",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
Cursor = Annotated[
|
||||
str | None,
|
||||
Field(
|
||||
description="Pagination cursor to continue from previous batch results",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class BatchToolParams(TypedDict):
|
||||
"""Parameters for the BatchTool.
|
||||
|
||||
Attributes:
|
||||
description: A short (3-5 word) description of the batch operation
|
||||
invocations: The list of tool invocations to execute (required -- you MUST provide at least one tool invocation)
|
||||
cursor: Optional pagination cursor
|
||||
"""
|
||||
|
||||
description: Description
|
||||
invocations: Invocations
|
||||
cursor: Cursor
|
||||
|
||||
|
||||
@final
|
||||
class BatchTool(BaseTool):
|
||||
"""Tool for executing multiple tools in a single request.
|
||||
|
||||
Executes a list of tool invocations in parallel when possible, or
|
||||
otherwise serially. Returns the collected results from all invocations.
|
||||
"""
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name.
|
||||
|
||||
Returns:
|
||||
Tool name
|
||||
"""
|
||||
return "batch"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description.
|
||||
|
||||
Returns:
|
||||
Tool description
|
||||
"""
|
||||
return """Batch execution tool that runs multiple tool invocations in a single request.
|
||||
|
||||
Tools are executed in parallel when possible, and otherwise serially.
|
||||
Takes a list of tool invocations (tool_name and input pairs).
|
||||
Returns the collected results from all invocations.
|
||||
Use this tool when you need to run multiple independent tool operations at once -- it is awesome for speeding up your workflow, reducing both context usage and latency.
|
||||
Each tool will respect its own permissions and validation rules.
|
||||
The tool's outputs are NOT shown to the user; to answer the user's query, you MUST send a message with the results after the tool call completes, otherwise the user will not see the results.
|
||||
|
||||
<batch_example>
|
||||
When dispatching multiple agents to find necessary information.
|
||||
batch(
|
||||
description="Update import statements across modules",
|
||||
invocations=[
|
||||
{tool_name: "dispatch_agent", input: {prompt: "Search for all instances of 'logger' configuration in /app/config directory"}},
|
||||
{tool_name: "dispatch_agent", input: {prompt: "Find all test files that reference 'UserService' in /app/tests"}},
|
||||
]
|
||||
)
|
||||
|
||||
Common scenarios for effective batching:
|
||||
1. Reading multiple related files in one operation
|
||||
2. Performing a series of simple mechanical changes
|
||||
3. Running multiple diagnostic commands
|
||||
4. Dispatch multiple agents to complete the task
|
||||
|
||||
To make a batch call, provide the following:
|
||||
1. description: A short (3-5 word) description of the batch operation
|
||||
2. invocations: List of invocation [{"tool_name": "...", "input": "..."}], tool_name: The name of the tool to invoke,newText: The input to pass to the tool
|
||||
|
||||
|
||||
Available tools in batch call:
|
||||
Tool: dispatch_agent,read,directory_tree,grep,grep_ast,run_command,notebook_read
|
||||
CLI Tools: claude,cc,codex,gemini,grok,openhands,oh,hanzo_dev,cline,aider
|
||||
AST/Symbols: ast,symbols (tree-sitter based code analysis)
|
||||
Not available: think,write,edit,multi_edit,notebook_edit
|
||||
"""
|
||||
|
||||
def __init__(self, tools: dict[str, BaseTool]) -> None:
|
||||
"""Initialize the batch tool.
|
||||
|
||||
Args:
|
||||
tools: Dictionary mapping tool names to tool instances
|
||||
"""
|
||||
self.tools = tools
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[BatchToolParams],
|
||||
) -> str:
|
||||
"""Execute the tool with the given parameters.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Tool result
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
description = params.get("description")
|
||||
invocations: list[dict[str, Any]] = params.get("invocations", list())
|
||||
|
||||
# Validate required parameters
|
||||
if not description:
|
||||
await tool_ctx.error(
|
||||
"Parameter 'description' is required but was None or empty"
|
||||
)
|
||||
return "Error: Parameter 'description' is required but was None or empty"
|
||||
|
||||
if not invocations:
|
||||
await tool_ctx.error(
|
||||
"Parameter 'invocations' is required but was None or empty"
|
||||
)
|
||||
return "Error: Parameter 'invocations' is required but was None or empty"
|
||||
|
||||
if not isinstance(invocations, list) or len(invocations) == 0:
|
||||
await tool_ctx.error("Parameter 'invocations' must be a non-empty list")
|
||||
return "Error: Parameter 'invocations' must be a non-empty list"
|
||||
|
||||
await tool_ctx.info(
|
||||
f"Executing batch operation: {description} ({len(invocations)} invocations)"
|
||||
)
|
||||
|
||||
# Execute all tool invocations in parallel
|
||||
tasks: list[asyncio.Future[dict[str, Any]]] = []
|
||||
invocation_map: dict[asyncio.Future[dict[str, Any]], dict[str, Any]] = (
|
||||
{}
|
||||
) # Map task Future to invocation
|
||||
|
||||
for i, invocation in enumerate(invocations):
|
||||
# Extract tool name and input from invocation
|
||||
tool_name: str = invocation.get("tool_name", "")
|
||||
tool_input: dict[str, Any] = invocation.get("input", {})
|
||||
|
||||
# Validate tool invocation
|
||||
if not tool_name:
|
||||
error_message = f"Tool name is required in invocation {i}"
|
||||
await tool_ctx.error(error_message)
|
||||
# Add direct result for this invocation
|
||||
tasks.append(asyncio.Future())
|
||||
tasks[-1].set_result(
|
||||
{"invocation": invocation, "result": f"Error: {error_message}"}
|
||||
)
|
||||
invocation_map[tasks[-1]] = invocation
|
||||
continue
|
||||
|
||||
# Check if the tool exists
|
||||
if tool_name not in self.tools:
|
||||
error_message = f"Tool '{tool_name}' not found"
|
||||
await tool_ctx.error(error_message)
|
||||
# Add direct result for this invocation
|
||||
tasks.append(asyncio.Future())
|
||||
tasks[-1].set_result(
|
||||
{"invocation": invocation, "result": f"Error: {error_message}"}
|
||||
)
|
||||
invocation_map[tasks[-1]] = invocation
|
||||
continue
|
||||
|
||||
# Create a task for this tool invocation
|
||||
try:
|
||||
tool = self.tools[tool_name]
|
||||
await tool_ctx.info(f"Creating task for tool: {tool_name}")
|
||||
|
||||
# Create coroutine for this tool execution
|
||||
async def execute_tool(
|
||||
tool_obj: BaseTool, tool_name: str, tool_input: dict[str, Any]
|
||||
):
|
||||
try:
|
||||
await tool_ctx.info(f"Executing tool: {tool_name}")
|
||||
result = await tool_obj.call(ctx, **tool_input)
|
||||
await tool_ctx.info(f"Tool '{tool_name}' execution completed")
|
||||
return {
|
||||
"invocation": {"tool_name": tool_name, "input": tool_input},
|
||||
"result": result,
|
||||
}
|
||||
except Exception as e:
|
||||
error_message = f"Error executing tool '{tool_name}': {str(e)}"
|
||||
await tool_ctx.error(error_message)
|
||||
return {
|
||||
"invocation": {"tool_name": tool_name, "input": tool_input},
|
||||
"result": f"Error: {error_message}",
|
||||
}
|
||||
|
||||
# Schedule the task
|
||||
task = asyncio.create_task(execute_tool(tool, tool_name, tool_input))
|
||||
tasks.append(task)
|
||||
invocation_map[task] = invocation
|
||||
except Exception as e:
|
||||
error_message = f"Error scheduling tool '{tool_name}': {str(e)}"
|
||||
await tool_ctx.error(error_message)
|
||||
# Add direct result for this invocation
|
||||
tasks.append(asyncio.Future())
|
||||
tasks[-1].set_result(
|
||||
{"invocation": invocation, "result": f"Error: {error_message}"}
|
||||
)
|
||||
invocation_map[tasks[-1]] = invocation
|
||||
|
||||
# Wait for all tasks to complete
|
||||
await tool_ctx.info(f"Waiting for {len(tasks)} tool executions to complete")
|
||||
results: list[dict[str, Any]] = []
|
||||
|
||||
# As tasks complete, collect their results
|
||||
for task in asyncio.as_completed(tasks):
|
||||
try:
|
||||
result = await task
|
||||
results.append(result)
|
||||
except Exception as e:
|
||||
invocation = invocation_map[task]
|
||||
tool_name: str = invocation.get("tool_name", "unknown")
|
||||
error_message = f"Unexpected error in tool '{tool_name}': {str(e)}"
|
||||
await tool_ctx.error(error_message)
|
||||
results.append(
|
||||
{"invocation": invocation, "result": f"Error: {error_message}"}
|
||||
)
|
||||
|
||||
# Extract cursor if provided
|
||||
cursor = params.get("cursor")
|
||||
cursor_offset = 0
|
||||
|
||||
# If cursor provided, we need to resume from where we left off
|
||||
if cursor:
|
||||
cursor_data = CursorData.from_cursor(cursor)
|
||||
if cursor_data and cursor_data.offset < len(results):
|
||||
# Skip already returned results
|
||||
cursor_offset = cursor_data.offset
|
||||
results = results[cursor_offset:]
|
||||
|
||||
# Format results
|
||||
formatted_results = []
|
||||
for i, result in enumerate(results):
|
||||
invocation = result["invocation"]
|
||||
tool_name = invocation.get("tool_name", "unknown")
|
||||
formatted_results.append(
|
||||
{
|
||||
"tool": tool_name,
|
||||
"result": result["result"],
|
||||
"index": i + cursor_offset,
|
||||
}
|
||||
)
|
||||
|
||||
# Create paginated response with token awareness
|
||||
paginated_response = create_paginated_response(
|
||||
formatted_results, cursor=cursor, use_token_limit=True
|
||||
)
|
||||
|
||||
# Convert paginated response to string format for MCP
|
||||
if isinstance(paginated_response, dict) and "items" in paginated_response:
|
||||
# Format the items as a readable string
|
||||
result_parts = []
|
||||
|
||||
# Add header
|
||||
result_parts.append(f"=== Batch operation: {description} ===")
|
||||
result_parts.append(f"Total invocations: {len(invocations)}")
|
||||
result_parts.append(
|
||||
f"Showing results: {len(paginated_response['items'])} of {len(results)}"
|
||||
)
|
||||
if paginated_response.get("hasMore"):
|
||||
result_parts.append(
|
||||
f"More results available - use cursor: {paginated_response.get('nextCursor')}"
|
||||
)
|
||||
result_parts.append("")
|
||||
|
||||
# Format each result
|
||||
for item in paginated_response["items"]:
|
||||
result_parts.append(f"### Result {item['index'] + 1}: {item['tool']}")
|
||||
result_content = item["result"]
|
||||
|
||||
# Add the result content - use multi-line code blocks for code outputs
|
||||
if isinstance(result_content, str) and "\n" in result_content:
|
||||
result_parts.append(f"```\n{result_content}\n```")
|
||||
else:
|
||||
result_parts.append(str(result_content))
|
||||
result_parts.append("")
|
||||
|
||||
# Join all parts
|
||||
formatted_output = "\n".join(result_parts)
|
||||
|
||||
# If there's a next cursor, we need to preserve it in the response
|
||||
# For now, append it as a note at the end
|
||||
if paginated_response.get("hasMore") and paginated_response.get(
|
||||
"nextCursor"
|
||||
):
|
||||
formatted_output += (
|
||||
f"\n\n[To continue, use cursor: {paginated_response['nextCursor']}]"
|
||||
)
|
||||
|
||||
await tool_ctx.info(
|
||||
f"Batch operation '{description}' completed with {len(paginated_response['items'])} results"
|
||||
f"{' (more available)' if paginated_response.get('hasMore') else ''}"
|
||||
)
|
||||
|
||||
return formatted_output
|
||||
else:
|
||||
# Fallback if pagination didn't work as expected
|
||||
return self._format_results(results)
|
||||
|
||||
def _format_results(self, results: list[dict[str, dict[str, Any]]]) -> str:
|
||||
"""Format the results from multiple tool invocations.
|
||||
|
||||
Args:
|
||||
results: List of tool invocation results
|
||||
|
||||
Returns:
|
||||
Formatted results string
|
||||
"""
|
||||
formatted_parts: list[str] = []
|
||||
for i, result in enumerate(results):
|
||||
invocation: dict[str, Any] = result["invocation"]
|
||||
tool_name: str = invocation.get("tool_name", "unknown")
|
||||
|
||||
# Add the result header
|
||||
formatted_parts.append(f"### Result {i + 1}: {tool_name}")
|
||||
|
||||
# Truncate individual results if they're too large
|
||||
result_content = result["result"]
|
||||
if len(result_content) > 50000: # If individual result > 50k chars
|
||||
result_content = truncate_response(
|
||||
result_content,
|
||||
max_tokens=5000, # Limit individual results to ~5k tokens
|
||||
truncation_message=f"\n\n[Result from {tool_name} truncated. Use the tool directly with pagination/filtering for full output.]",
|
||||
)
|
||||
|
||||
# Add the result content - use multi-line code blocks for code outputs
|
||||
if "\n" in result_content:
|
||||
formatted_parts.append(f"```\n{result_content}\n```")
|
||||
else:
|
||||
formatted_parts.append(result_content)
|
||||
# Add a separator
|
||||
formatted_parts.append("")
|
||||
|
||||
return "\n".join(formatted_parts)
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this batch tool with the MCP server.
|
||||
|
||||
Creates a wrapper function with explicitly defined parameters that match
|
||||
the tool's parameter schema and registers it with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
"""
|
||||
tool_self = self # Create a reference to self for use in the closure
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def batch(
|
||||
description: Description,
|
||||
invocations: Invocations,
|
||||
cursor: Cursor,
|
||||
ctx: MCPContext,
|
||||
) -> str:
|
||||
return await tool_self.call(
|
||||
ctx, description=description, invocations=invocations, cursor=cursor
|
||||
)
|
||||
@@ -1,432 +0,0 @@
|
||||
"""Configuration management tool for dynamic settings updates."""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, Unpack, Optional, TypedDict, final
|
||||
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.config.settings import (
|
||||
ProjectConfig,
|
||||
MCPServerConfig,
|
||||
load_settings,
|
||||
save_settings,
|
||||
detect_project_from_path,
|
||||
)
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
|
||||
class ConfigToolParams(TypedDict, total=False):
|
||||
"""Parameters for configuration management operations."""
|
||||
|
||||
action: str # get, set, add_server, remove_server, add_project, etc.
|
||||
scope: Optional[str] # global, project, current
|
||||
setting_path: Optional[str] # dot-notation path like "agent.enabled"
|
||||
value: Optional[Any] # new value to set
|
||||
server_name: Optional[str] # for MCP server operations
|
||||
server_config: Optional[Dict[str, Any]] # for adding servers
|
||||
project_name: Optional[str] # for project operations
|
||||
project_path: Optional[str] # for project path detection
|
||||
|
||||
|
||||
@final
|
||||
class ConfigTool(BaseTool):
|
||||
"""Tool for managing Hanzo AI configuration dynamically."""
|
||||
|
||||
def __init__(self, permission_manager: PermissionManager):
|
||||
"""Initialize the configuration tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "config"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Dynamically manage Hanzo AI configuration settings through conversation.
|
||||
|
||||
Can get/set global settings, project-specific settings, manage MCP servers, configure tools,
|
||||
and handle project workflows. Supports dot-notation for nested settings like 'agent.enabled'.
|
||||
|
||||
Perfect for AI-driven configuration where users can say things like:
|
||||
- "Enable the agent tool for this project"
|
||||
- "Add a new MCP server for file operations"
|
||||
- "Disable write tools globally but enable them for the current project"
|
||||
- "Show me the current project configuration"
|
||||
|
||||
Automatically detects projects based on LLM.md files and manages .hanzo/ directories."""
|
||||
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[ConfigToolParams],
|
||||
) -> str:
|
||||
"""Manage configuration settings.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Configuration operation result
|
||||
"""
|
||||
action = params.get("action", "get")
|
||||
scope = params.get("scope", "global")
|
||||
setting_path = params.get("setting_path")
|
||||
value = params.get("value")
|
||||
server_name = params.get("server_name")
|
||||
server_config = params.get("server_config")
|
||||
project_name = params.get("project_name")
|
||||
project_path = params.get("project_path")
|
||||
|
||||
try:
|
||||
if action == "get":
|
||||
return await self._get_config(
|
||||
scope, setting_path, project_name, project_path
|
||||
)
|
||||
elif action == "set":
|
||||
return await self._set_config(
|
||||
scope, setting_path, value, project_name, project_path
|
||||
)
|
||||
elif action == "add_server":
|
||||
return await self._add_mcp_server(
|
||||
server_name, server_config, scope, project_name
|
||||
)
|
||||
elif action == "remove_server":
|
||||
return await self._remove_mcp_server(server_name, scope, project_name)
|
||||
elif action == "enable_server":
|
||||
return await self._enable_mcp_server(server_name, scope, project_name)
|
||||
elif action == "disable_server":
|
||||
return await self._disable_mcp_server(server_name, scope, project_name)
|
||||
elif action == "trust_server":
|
||||
return await self._trust_mcp_server(server_name)
|
||||
elif action == "add_project":
|
||||
return await self._add_project(project_name, project_path)
|
||||
elif action == "set_current_project":
|
||||
return await self._set_current_project(project_name, project_path)
|
||||
elif action == "list_servers":
|
||||
return await self._list_mcp_servers(scope, project_name)
|
||||
elif action == "list_projects":
|
||||
return await self._list_projects()
|
||||
elif action == "detect_project":
|
||||
return await self._detect_project(project_path)
|
||||
else:
|
||||
return f"Error: Unknown action '{action}'. Available actions: get, set, add_server, remove_server, enable_server, disable_server, trust_server, add_project, set_current_project, list_servers, list_projects, detect_project"
|
||||
|
||||
except Exception as e:
|
||||
return f"Error managing configuration: {str(e)}"
|
||||
|
||||
async def _get_config(
|
||||
self,
|
||||
scope: str,
|
||||
setting_path: Optional[str],
|
||||
project_name: Optional[str],
|
||||
project_path: Optional[str],
|
||||
) -> str:
|
||||
"""Get configuration value(s)."""
|
||||
# Load appropriate settings
|
||||
if scope == "project" or project_name or project_path:
|
||||
if project_path:
|
||||
project_info = detect_project_from_path(project_path)
|
||||
if project_info:
|
||||
settings = load_settings(project_info["root_path"])
|
||||
else:
|
||||
return f"No project detected at path: {project_path}"
|
||||
else:
|
||||
settings = load_settings()
|
||||
else:
|
||||
settings = load_settings()
|
||||
|
||||
if not setting_path:
|
||||
# Return full config summary
|
||||
if scope == "project" and settings.current_project:
|
||||
project = settings.get_current_project()
|
||||
if project:
|
||||
return f"Current Project Configuration ({project.name}):\\n{json.dumps(project.__dict__, indent=2)}"
|
||||
|
||||
# Return global config summary
|
||||
summary = {
|
||||
"server": settings.server.__dict__,
|
||||
"enabled_tools": settings.get_enabled_tools(),
|
||||
"disabled_tools": settings.get_disabled_tools(),
|
||||
"agent": settings.agent.__dict__,
|
||||
"vector_store": settings.vector_store.__dict__,
|
||||
"hub_enabled": settings.hub_enabled,
|
||||
"mcp_servers": {
|
||||
name: server.__dict__
|
||||
for name, server in settings.mcp_servers.items()
|
||||
},
|
||||
"current_project": settings.current_project,
|
||||
"projects": list(settings.projects.keys()),
|
||||
}
|
||||
return f"Configuration Summary:\\n{json.dumps(summary, indent=2)}"
|
||||
|
||||
# Get specific setting
|
||||
value = self._get_nested_value(settings.__dict__, setting_path)
|
||||
if value is not None:
|
||||
return f"{setting_path}: {json.dumps(value, indent=2)}"
|
||||
else:
|
||||
return f"Setting '{setting_path}' not found"
|
||||
|
||||
async def _set_config(
|
||||
self,
|
||||
scope: str,
|
||||
setting_path: Optional[str],
|
||||
value: Any,
|
||||
project_name: Optional[str],
|
||||
project_path: Optional[str],
|
||||
) -> str:
|
||||
"""Set configuration value."""
|
||||
if not setting_path:
|
||||
return "Error: setting_path is required for set action"
|
||||
|
||||
# Load settings
|
||||
project_dir = None
|
||||
if scope == "project" or project_name or project_path:
|
||||
if project_path:
|
||||
project_info = detect_project_from_path(project_path)
|
||||
if project_info:
|
||||
project_dir = project_info["root_path"]
|
||||
else:
|
||||
return f"No project detected at path: {project_path}"
|
||||
|
||||
settings = load_settings(project_dir)
|
||||
|
||||
# Set the value
|
||||
if self._set_nested_value(settings.__dict__, setting_path, value):
|
||||
# Save settings
|
||||
if scope == "project" or project_dir:
|
||||
save_settings(settings, global_config=False)
|
||||
else:
|
||||
save_settings(settings, global_config=True)
|
||||
return f"Successfully set {setting_path} = {json.dumps(value)}"
|
||||
else:
|
||||
return f"Error: Could not set '{setting_path}'"
|
||||
|
||||
async def _add_mcp_server(
|
||||
self,
|
||||
server_name: Optional[str],
|
||||
server_config: Optional[Dict[str, Any]],
|
||||
scope: str,
|
||||
project_name: Optional[str],
|
||||
) -> str:
|
||||
"""Add a new MCP server."""
|
||||
if not server_name or not server_config:
|
||||
return "Error: server_name and server_config are required"
|
||||
|
||||
settings = load_settings()
|
||||
|
||||
# Create server config
|
||||
mcp_server = MCPServerConfig(
|
||||
name=server_name,
|
||||
command=server_config.get("command", ""),
|
||||
args=server_config.get("args", []),
|
||||
enabled=server_config.get("enabled", True),
|
||||
trusted=server_config.get("trusted", False),
|
||||
description=server_config.get("description", ""),
|
||||
capabilities=server_config.get("capabilities", []),
|
||||
)
|
||||
|
||||
if settings.add_mcp_server(mcp_server):
|
||||
save_settings(settings)
|
||||
return f"Successfully added MCP server '{server_name}'"
|
||||
else:
|
||||
return f"Error: MCP server '{server_name}' already exists"
|
||||
|
||||
async def _remove_mcp_server(
|
||||
self, server_name: Optional[str], scope: str, project_name: Optional[str]
|
||||
) -> str:
|
||||
"""Remove an MCP server."""
|
||||
if not server_name:
|
||||
return "Error: server_name is required"
|
||||
|
||||
settings = load_settings()
|
||||
|
||||
if settings.remove_mcp_server(server_name):
|
||||
save_settings(settings)
|
||||
return f"Successfully removed MCP server '{server_name}'"
|
||||
else:
|
||||
return f"Error: MCP server '{server_name}' not found"
|
||||
|
||||
async def _enable_mcp_server(
|
||||
self, server_name: Optional[str], scope: str, project_name: Optional[str]
|
||||
) -> str:
|
||||
"""Enable an MCP server."""
|
||||
if not server_name:
|
||||
return "Error: server_name is required"
|
||||
|
||||
settings = load_settings()
|
||||
|
||||
if settings.enable_mcp_server(server_name):
|
||||
save_settings(settings)
|
||||
return f"Successfully enabled MCP server '{server_name}'"
|
||||
else:
|
||||
return f"Error: MCP server '{server_name}' not found"
|
||||
|
||||
async def _disable_mcp_server(
|
||||
self, server_name: Optional[str], scope: str, project_name: Optional[str]
|
||||
) -> str:
|
||||
"""Disable an MCP server."""
|
||||
if not server_name:
|
||||
return "Error: server_name is required"
|
||||
|
||||
settings = load_settings()
|
||||
|
||||
if settings.disable_mcp_server(server_name):
|
||||
save_settings(settings)
|
||||
return f"Successfully disabled MCP server '{server_name}'"
|
||||
else:
|
||||
return f"Error: MCP server '{server_name}' not found"
|
||||
|
||||
async def _trust_mcp_server(self, server_name: Optional[str]) -> str:
|
||||
"""Trust an MCP server."""
|
||||
if not server_name:
|
||||
return "Error: server_name is required"
|
||||
|
||||
settings = load_settings()
|
||||
|
||||
if settings.trust_mcp_server(server_name):
|
||||
save_settings(settings)
|
||||
return f"Successfully trusted MCP server '{server_name}'"
|
||||
else:
|
||||
return f"Error: MCP server '{server_name}' not found"
|
||||
|
||||
async def _add_project(
|
||||
self, project_name: Optional[str], project_path: Optional[str]
|
||||
) -> str:
|
||||
"""Add a project configuration."""
|
||||
if not project_path:
|
||||
return "Error: project_path is required"
|
||||
|
||||
# Detect or create project
|
||||
project_info = detect_project_from_path(project_path)
|
||||
if not project_info:
|
||||
return f"No LLM.md found in project path: {project_path}"
|
||||
|
||||
if not project_name:
|
||||
project_name = project_info["name"]
|
||||
|
||||
settings = load_settings()
|
||||
|
||||
project_config = ProjectConfig(
|
||||
name=project_name,
|
||||
root_path=project_info["root_path"],
|
||||
)
|
||||
|
||||
if settings.add_project(project_config):
|
||||
save_settings(settings)
|
||||
return f"Successfully added project '{project_name}' at {project_info['root_path']}"
|
||||
else:
|
||||
return f"Error: Project '{project_name}' already exists"
|
||||
|
||||
async def _set_current_project(
|
||||
self, project_name: Optional[str], project_path: Optional[str]
|
||||
) -> str:
|
||||
"""Set the current active project."""
|
||||
settings = load_settings()
|
||||
|
||||
if project_path:
|
||||
project_info = detect_project_from_path(project_path)
|
||||
if project_info:
|
||||
project_name = project_info["name"]
|
||||
# Auto-add project if not exists
|
||||
if project_name not in settings.projects:
|
||||
await self._add_project(project_name, project_path)
|
||||
settings = load_settings() # Reload after adding
|
||||
|
||||
if not project_name:
|
||||
return "Error: project_name or project_path is required"
|
||||
|
||||
if settings.set_current_project(project_name):
|
||||
save_settings(settings)
|
||||
return f"Successfully set current project to '{project_name}'"
|
||||
else:
|
||||
return f"Error: Project '{project_name}' not found"
|
||||
|
||||
async def _list_mcp_servers(self, scope: str, project_name: Optional[str]) -> str:
|
||||
"""List MCP servers."""
|
||||
settings = load_settings()
|
||||
|
||||
if not settings.mcp_servers:
|
||||
return "No MCP servers configured"
|
||||
|
||||
servers_info = []
|
||||
for name, server in settings.mcp_servers.items():
|
||||
status = "enabled" if server.enabled else "disabled"
|
||||
trust = "trusted" if server.trusted else "untrusted"
|
||||
servers_info.append(f"- {name}: {server.command} ({status}, {trust})")
|
||||
|
||||
return f"MCP Servers:\\n" + "\\n".join(servers_info)
|
||||
|
||||
async def _list_projects(self) -> str:
|
||||
"""List projects."""
|
||||
settings = load_settings()
|
||||
|
||||
if not settings.projects:
|
||||
return "No projects configured"
|
||||
|
||||
projects_info = []
|
||||
for name, project in settings.projects.items():
|
||||
current = " (current)" if name == settings.current_project else ""
|
||||
projects_info.append(f"- {name}: {project.root_path}{current}")
|
||||
|
||||
return f"Projects:\\n" + "\\n".join(projects_info)
|
||||
|
||||
async def _detect_project(self, project_path: Optional[str]) -> str:
|
||||
"""Detect project from path."""
|
||||
if not project_path:
|
||||
import os
|
||||
|
||||
project_path = os.getcwd()
|
||||
|
||||
project_info = detect_project_from_path(project_path)
|
||||
if project_info:
|
||||
return f"Project detected:\\n{json.dumps(project_info, indent=2)}"
|
||||
else:
|
||||
return f"No project detected at path: {project_path}"
|
||||
|
||||
def _get_nested_value(self, obj: Dict[str, Any], path: str) -> Any:
|
||||
"""Get nested value using dot notation."""
|
||||
keys = path.split(".")
|
||||
current = obj
|
||||
|
||||
for key in keys:
|
||||
if isinstance(current, dict) and key in current:
|
||||
current = current[key]
|
||||
elif hasattr(current, key):
|
||||
current = getattr(current, key)
|
||||
else:
|
||||
return None
|
||||
|
||||
return current
|
||||
|
||||
def _set_nested_value(self, obj: Dict[str, Any], path: str, value: Any) -> bool:
|
||||
"""Set nested value using dot notation."""
|
||||
keys = path.split(".")
|
||||
current = obj
|
||||
|
||||
# Navigate to parent
|
||||
for key in keys[:-1]:
|
||||
if isinstance(current, dict) and key in current:
|
||||
current = current[key]
|
||||
elif hasattr(current, key):
|
||||
current = getattr(current, key)
|
||||
else:
|
||||
return False
|
||||
|
||||
# Set final value
|
||||
final_key = keys[-1]
|
||||
if isinstance(current, dict):
|
||||
current[final_key] = value
|
||||
return True
|
||||
elif hasattr(current, final_key):
|
||||
setattr(current, final_key, value)
|
||||
return True
|
||||
|
||||
return False
|
||||
@@ -1,180 +0,0 @@
|
||||
"""Enhanced Context for Hanzo AI tools.
|
||||
|
||||
This module provides an enhanced Context class that wraps the MCP Context
|
||||
and adds additional functionality specific to Hanzo tools.
|
||||
"""
|
||||
|
||||
from typing import ClassVar, final
|
||||
from collections.abc import Iterable
|
||||
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
from mcp.server.lowlevel.helper_types import ReadResourceContents
|
||||
|
||||
|
||||
@final
|
||||
class ToolContext:
|
||||
"""Enhanced context for Hanzo AI tools.
|
||||
|
||||
This class wraps the MCP Context and adds additional functionality
|
||||
for tracking tool execution, progress reporting, and resource access.
|
||||
"""
|
||||
|
||||
# Track all active contexts for debugging
|
||||
_active_contexts: ClassVar[set["ToolContext"]] = set()
|
||||
|
||||
def __init__(self, mcp_context: MCPContext) -> None:
|
||||
"""Initialize the tool context.
|
||||
|
||||
Args:
|
||||
mcp_context: The underlying MCP Context
|
||||
"""
|
||||
self._mcp_context: MCPContext = mcp_context
|
||||
self._tool_name: str | None = None
|
||||
self._execution_id: str | None = None
|
||||
|
||||
# Add to active contexts
|
||||
ToolContext._active_contexts.add(self)
|
||||
|
||||
def __del__(self) -> None:
|
||||
"""Clean up when the context is destroyed."""
|
||||
# Remove from active contexts
|
||||
ToolContext._active_contexts.discard(self)
|
||||
|
||||
@property
|
||||
def mcp_context(self) -> MCPContext:
|
||||
"""Get the underlying MCP Context.
|
||||
|
||||
Returns:
|
||||
The MCP Context
|
||||
"""
|
||||
return self._mcp_context
|
||||
|
||||
@property
|
||||
def request_id(self) -> str:
|
||||
"""Get the request ID from the MCP context.
|
||||
|
||||
Returns:
|
||||
The request ID
|
||||
"""
|
||||
return self._mcp_context.request_id
|
||||
|
||||
@property
|
||||
def client_id(self) -> str | None:
|
||||
"""Get the client ID from the MCP context.
|
||||
|
||||
Returns:
|
||||
The client ID
|
||||
"""
|
||||
return self._mcp_context.client_id
|
||||
|
||||
async def set_tool_info(
|
||||
self, tool_name: str, execution_id: str | None = None
|
||||
) -> None:
|
||||
"""Set information about the currently executing tool.
|
||||
|
||||
Args:
|
||||
tool_name: The name of the tool being executed
|
||||
execution_id: Optional unique execution ID
|
||||
"""
|
||||
self._tool_name = tool_name
|
||||
self._execution_id = execution_id
|
||||
|
||||
async def info(self, message: str) -> None:
|
||||
"""Log an informational message.
|
||||
|
||||
Args:
|
||||
message: The message to log
|
||||
"""
|
||||
try:
|
||||
await self._mcp_context.info(self._format_message(message))
|
||||
except Exception:
|
||||
# Silently ignore errors when client has disconnected
|
||||
pass
|
||||
|
||||
async def debug(self, message: str) -> None:
|
||||
"""Log a debug message.
|
||||
|
||||
Args:
|
||||
message: The message to log
|
||||
"""
|
||||
try:
|
||||
await self._mcp_context.debug(self._format_message(message))
|
||||
except Exception:
|
||||
# Silently ignore errors when client has disconnected
|
||||
pass
|
||||
|
||||
async def warning(self, message: str) -> None:
|
||||
"""Log a warning message.
|
||||
|
||||
Args:
|
||||
message: The message to log
|
||||
"""
|
||||
try:
|
||||
await self._mcp_context.warning(self._format_message(message))
|
||||
except Exception:
|
||||
# Silently ignore errors when client has disconnected
|
||||
pass
|
||||
|
||||
async def error(self, message: str) -> None:
|
||||
"""Log an error message.
|
||||
|
||||
Args:
|
||||
message: The message to log
|
||||
"""
|
||||
try:
|
||||
await self._mcp_context.error(self._format_message(message))
|
||||
except Exception:
|
||||
# Silently ignore errors when client has disconnected
|
||||
pass
|
||||
|
||||
def _format_message(self, message: str) -> str:
|
||||
"""Format a message with tool information if available.
|
||||
|
||||
Args:
|
||||
message: The original message
|
||||
|
||||
Returns:
|
||||
The formatted message
|
||||
"""
|
||||
if self._tool_name:
|
||||
if self._execution_id:
|
||||
return f"[{self._tool_name}:{self._execution_id}] {message}"
|
||||
return f"[{self._tool_name}] {message}"
|
||||
return message
|
||||
|
||||
async def report_progress(self, current: int, total: int) -> None:
|
||||
"""Report progress to the client.
|
||||
|
||||
Args:
|
||||
current: Current progress value
|
||||
total: Total progress value
|
||||
"""
|
||||
try:
|
||||
await self._mcp_context.report_progress(current, total)
|
||||
except Exception:
|
||||
# Silently ignore errors when client has disconnected
|
||||
pass
|
||||
|
||||
async def read_resource(self, uri: str) -> Iterable[ReadResourceContents]:
|
||||
"""Read a resource via the MCP protocol.
|
||||
|
||||
Args:
|
||||
uri: The resource URI
|
||||
|
||||
Returns:
|
||||
A tuple of (content, mime_type)
|
||||
"""
|
||||
return await self._mcp_context.read_resource(uri)
|
||||
|
||||
|
||||
# Factory function to create a ToolContext from an MCP Context
|
||||
def create_tool_context(mcp_context: MCPContext) -> ToolContext:
|
||||
"""Create a ToolContext from an MCP Context.
|
||||
|
||||
Args:
|
||||
mcp_context: The MCP Context
|
||||
|
||||
Returns:
|
||||
A new ToolContext
|
||||
"""
|
||||
return ToolContext(mcp_context)
|
||||
@@ -1,28 +0,0 @@
|
||||
"""Context handling fix for MCP tools.
|
||||
|
||||
This module provides backward compatibility by re-exporting the
|
||||
context normalization utilities from the decorators module.
|
||||
|
||||
DEPRECATED: Use hanzo_mcp.tools.common.decorators directly.
|
||||
"""
|
||||
|
||||
# Re-export for backward compatibility
|
||||
from hanzo_mcp.tools.common.decorators import (
|
||||
MockContext,
|
||||
_is_valid_context as is_valid_context,
|
||||
with_context_normalization,
|
||||
)
|
||||
|
||||
|
||||
# Backward compatibility function
|
||||
def normalize_context(ctx):
|
||||
"""Normalize context - backward compatibility wrapper.
|
||||
|
||||
DEPRECATED: Use decorators.with_context_normalization instead.
|
||||
"""
|
||||
if is_valid_context(ctx):
|
||||
return ctx
|
||||
return MockContext()
|
||||
|
||||
|
||||
__all__ = ["MockContext", "normalize_context", "with_context_normalization"]
|
||||
@@ -1,192 +0,0 @@
|
||||
"""Critic tool implementation.
|
||||
|
||||
This module provides the CriticTool for Claude to engage in critical analysis and code review.
|
||||
"""
|
||||
|
||||
from typing import Unpack, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
|
||||
Analysis = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="The critical analysis to perform - code review, error detection, or improvement suggestions",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class CriticToolParams(TypedDict):
|
||||
"""Parameters for the CriticTool.
|
||||
|
||||
Attributes:
|
||||
analysis: The critical analysis to perform
|
||||
"""
|
||||
|
||||
analysis: Analysis
|
||||
|
||||
|
||||
@final
|
||||
class CriticTool(BaseTool):
|
||||
"""Tool for Claude to engage in critical analysis and play devil's advocate."""
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name.
|
||||
|
||||
Returns:
|
||||
Tool name
|
||||
"""
|
||||
return "critic"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description.
|
||||
|
||||
Returns:
|
||||
Tool description
|
||||
"""
|
||||
return """Use this tool to perform critical analysis, play devil's advocate, and ensure high standards.
|
||||
This tool forces a critical thinking mode that reviews all code for errors, improvements, and edge cases.
|
||||
It ensures tests are run, tests pass, and maintains high quality standards.
|
||||
|
||||
This is your inner critic that:
|
||||
- Always questions assumptions
|
||||
- Looks for potential bugs and edge cases
|
||||
- Ensures proper error handling
|
||||
- Verifies test coverage
|
||||
- Checks for performance issues
|
||||
- Reviews security implications
|
||||
- Suggests improvements and refactoring
|
||||
- Ensures code follows best practices
|
||||
- Questions design decisions
|
||||
- Looks for missing documentation
|
||||
|
||||
Common use cases:
|
||||
1. Before finalizing any code changes - review for bugs, edge cases, and improvements
|
||||
2. After implementing a feature - critically analyze if it truly solves the problem
|
||||
3. When tests pass too easily - question if tests are comprehensive enough
|
||||
4. Before marking a task complete - ensure all quality standards are met
|
||||
5. When something seems too simple - look for hidden complexity or missing requirements
|
||||
6. After fixing a bug - analyze if the fix addresses root cause or just symptoms
|
||||
|
||||
<critic_example>
|
||||
Code Review Analysis:
|
||||
- Implementation Issues:
|
||||
* No error handling for network failures in API calls
|
||||
* Missing validation for user input boundaries
|
||||
* Race condition possible in concurrent updates
|
||||
* Memory leak potential in event listener registration
|
||||
|
||||
- Test Coverage Gaps:
|
||||
* No tests for error scenarios
|
||||
* Missing edge case: empty array input
|
||||
* No performance benchmarks for large datasets
|
||||
* Integration tests don't cover authentication failures
|
||||
|
||||
- Security Concerns:
|
||||
* SQL injection vulnerability in query construction
|
||||
* Missing rate limiting on public endpoints
|
||||
* Sensitive data logged in debug mode
|
||||
|
||||
- Performance Issues:
|
||||
* O(n²) algorithm where O(n log n) is possible
|
||||
* Database queries in a loop (N+1 problem)
|
||||
* No caching for expensive computations
|
||||
|
||||
- Code Quality:
|
||||
* Functions too long and doing multiple things
|
||||
* Inconsistent naming conventions
|
||||
* Missing type annotations
|
||||
* No documentation for complex algorithms
|
||||
|
||||
- Design Flaws:
|
||||
* Tight coupling between modules
|
||||
* Hard-coded configuration values
|
||||
* No abstraction for external dependencies
|
||||
* Violates single responsibility principle
|
||||
|
||||
Recommendations:
|
||||
1. Add comprehensive error handling with retry logic
|
||||
2. Implement input validation with clear error messages
|
||||
3. Use database transactions to prevent race conditions
|
||||
4. Add memory cleanup in component unmount
|
||||
5. Parameterize SQL queries to prevent injection
|
||||
6. Implement rate limiting middleware
|
||||
7. Use environment variables for sensitive config
|
||||
8. Refactor algorithm to use sorting approach
|
||||
9. Batch database queries
|
||||
10. Add memoization for expensive calculations
|
||||
</critic_example>"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialize the critic tool."""
|
||||
pass
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[CriticToolParams],
|
||||
) -> str:
|
||||
"""Execute the tool with the given parameters.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Tool result
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
analysis = params.get("analysis")
|
||||
|
||||
# Validate required analysis parameter
|
||||
if not analysis:
|
||||
await tool_ctx.error(
|
||||
"Parameter 'analysis' is required but was None or empty"
|
||||
)
|
||||
return "Error: Parameter 'analysis' is required but was None or empty"
|
||||
|
||||
if analysis.strip() == "":
|
||||
await tool_ctx.error("Parameter 'analysis' cannot be empty")
|
||||
return "Error: Parameter 'analysis' cannot be empty"
|
||||
|
||||
# Log the critical analysis
|
||||
await tool_ctx.info("Critical analysis recorded")
|
||||
|
||||
# Return confirmation with reminder to act on the analysis
|
||||
return """Critical analysis complete. Remember to:
|
||||
1. Address all identified issues before proceeding
|
||||
2. Run comprehensive tests to verify fixes
|
||||
3. Ensure all tests pass with proper coverage
|
||||
4. Document any design decisions or trade-offs
|
||||
5. Consider the analysis points in your implementation
|
||||
|
||||
Continue with improvements based on this critical review."""
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this critic tool with the MCP server.
|
||||
|
||||
Creates a wrapper function with explicitly defined parameters that match
|
||||
the tool's parameter schema and registers it with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
"""
|
||||
tool_self = self # Create a reference to self for use in the closure
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def critic(analysis: Analysis, ctx: MCPContext) -> str:
|
||||
return await tool_self.call(ctx, analysis=analysis)
|
||||
@@ -1,210 +0,0 @@
|
||||
"""Decorators for MCP tools.
|
||||
|
||||
This module provides decorators that handle common cross-cutting concerns
|
||||
for MCP tools, such as context normalization and error handling.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
import functools
|
||||
from typing import Any, TypeVar, Callable, cast
|
||||
|
||||
F = TypeVar("F", bound=Callable[..., Any])
|
||||
|
||||
|
||||
class MockContext:
|
||||
"""Mock context for when no real context is available.
|
||||
|
||||
This is used when tools are called externally through the MCP protocol
|
||||
and the Context parameter is not properly serialized.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.request_id = "external-request"
|
||||
self.client_id = "external-client"
|
||||
|
||||
async def info(self, message: str) -> None:
|
||||
"""Mock info logging - no-op for external calls."""
|
||||
pass
|
||||
|
||||
async def debug(self, message: str) -> None:
|
||||
"""Mock debug logging - no-op for external calls."""
|
||||
pass
|
||||
|
||||
async def warning(self, message: str) -> None:
|
||||
"""Mock warning logging - no-op for external calls."""
|
||||
pass
|
||||
|
||||
async def error(self, message: str) -> None:
|
||||
"""Mock error logging - no-op for external calls."""
|
||||
pass
|
||||
|
||||
async def report_progress(self, current: int, total: int) -> None:
|
||||
"""Mock progress reporting - no-op for external calls."""
|
||||
pass
|
||||
|
||||
async def read_resource(self, uri: str) -> Any:
|
||||
"""Mock resource reading - returns empty result."""
|
||||
return []
|
||||
|
||||
|
||||
def with_context_normalization(func: F) -> F:
|
||||
"""Decorator that normalizes the context parameter for MCP tools.
|
||||
|
||||
This decorator intercepts the ctx parameter and ensures it's a valid
|
||||
MCPContext object, even when called externally where it might be
|
||||
passed as a string, dict, or None.
|
||||
|
||||
Usage:
|
||||
@server.tool()
|
||||
@with_context_normalization
|
||||
async def my_tool(ctx: MCPContext, param: str) -> str:
|
||||
# ctx is guaranteed to be a valid context object
|
||||
await ctx.info("Processing...")
|
||||
return "result"
|
||||
|
||||
Args:
|
||||
func: The async function to decorate
|
||||
|
||||
Returns:
|
||||
The decorated function with context normalization
|
||||
"""
|
||||
|
||||
@functools.wraps(func)
|
||||
async def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
# Get function signature to find ctx parameter
|
||||
sig = inspect.signature(func)
|
||||
params = list(sig.parameters.keys())
|
||||
|
||||
# Handle ctx in kwargs
|
||||
if "ctx" in kwargs:
|
||||
ctx_value = kwargs["ctx"]
|
||||
if not _is_valid_context(ctx_value):
|
||||
kwargs["ctx"] = MockContext()
|
||||
|
||||
# Handle ctx in args (positional)
|
||||
elif "ctx" in params:
|
||||
ctx_index = params.index("ctx")
|
||||
if ctx_index < len(args):
|
||||
ctx_value = args[ctx_index]
|
||||
if not _is_valid_context(ctx_value):
|
||||
args_list = list(args)
|
||||
args_list[ctx_index] = MockContext()
|
||||
args = tuple(args_list)
|
||||
|
||||
# Call the original function
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
return cast(F, wrapper)
|
||||
|
||||
|
||||
def _is_valid_context(ctx: Any) -> bool:
|
||||
"""Check if an object is a valid MCPContext.
|
||||
|
||||
Args:
|
||||
ctx: The object to check
|
||||
|
||||
Returns:
|
||||
True if ctx is a valid context object
|
||||
"""
|
||||
# Check for required context methods
|
||||
return (
|
||||
hasattr(ctx, "info")
|
||||
and hasattr(ctx, "debug")
|
||||
and hasattr(ctx, "warning")
|
||||
and hasattr(ctx, "error")
|
||||
and hasattr(ctx, "report_progress")
|
||||
and
|
||||
# Ensure they're callable
|
||||
callable(getattr(ctx, "info", None))
|
||||
and callable(getattr(ctx, "debug", None))
|
||||
)
|
||||
|
||||
|
||||
def mcp_tool(
|
||||
server: Any, name: str | None = None, description: str | None = None
|
||||
) -> Callable[[F], F]:
|
||||
"""Enhanced MCP tool decorator that includes context normalization.
|
||||
|
||||
This decorator combines the standard MCP tool registration with
|
||||
automatic context normalization, providing a single-point solution
|
||||
for all tools.
|
||||
|
||||
Usage:
|
||||
@mcp_tool(server, name="my_tool", description="Does something")
|
||||
async def my_tool(ctx: MCPContext, param: str) -> str:
|
||||
await ctx.info("Processing...")
|
||||
return "result"
|
||||
|
||||
Args:
|
||||
server: The MCP server instance
|
||||
name: Optional tool name (defaults to function name)
|
||||
description: Optional tool description
|
||||
|
||||
Returns:
|
||||
Decorator function
|
||||
"""
|
||||
|
||||
def decorator(func: F) -> F:
|
||||
# Apply context normalization first
|
||||
normalized_func = with_context_normalization(func)
|
||||
|
||||
# Then apply the server's tool decorator
|
||||
if hasattr(server, "tool"):
|
||||
# Use the server's tool decorator
|
||||
server_decorator = server.tool(name=name, description=description)
|
||||
return server_decorator(normalized_func)
|
||||
else:
|
||||
# Fallback if server doesn't have tool method
|
||||
return normalized_func
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def create_tool_handler(server: Any, tool: Any) -> Callable[[], None]:
|
||||
"""Create a standardized tool registration handler.
|
||||
|
||||
This function creates a registration method that automatically applies
|
||||
context normalization to any tool handler registered with the server.
|
||||
|
||||
Usage:
|
||||
class MyTool(BaseTool):
|
||||
def register(self, mcp_server):
|
||||
register = create_tool_handler(mcp_server, self)
|
||||
register()
|
||||
|
||||
Args:
|
||||
server: The MCP server instance
|
||||
tool: The tool instance with name, description, and handler
|
||||
|
||||
Returns:
|
||||
A function that registers the tool with context normalization
|
||||
"""
|
||||
|
||||
def register_with_normalization():
|
||||
# Get the original register method
|
||||
original_register = tool.__class__.register
|
||||
|
||||
# Temporarily replace server.tool to wrap with normalization
|
||||
original_tool_decorator = server.tool
|
||||
|
||||
def normalized_tool_decorator(name=None, description=None):
|
||||
def decorator(func):
|
||||
# Apply context normalization
|
||||
normalized = with_context_normalization(func)
|
||||
# Apply original decorator
|
||||
return original_tool_decorator(name=name, description=description)(
|
||||
normalized
|
||||
)
|
||||
|
||||
return decorator
|
||||
|
||||
# Monkey-patch temporarily
|
||||
server.tool = normalized_tool_decorator
|
||||
try:
|
||||
# Call original register
|
||||
original_register(tool, server)
|
||||
finally:
|
||||
# Restore original
|
||||
server.tool = original_tool_decorator
|
||||
|
||||
return register_with_normalization
|
||||
@@ -1,103 +0,0 @@
|
||||
"""Enhanced base classes for MCP tools with automatic context handling.
|
||||
|
||||
This module provides enhanced base classes that automatically handle
|
||||
context normalization and other cross-cutting concerns, ensuring
|
||||
consistent behavior across all tools.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, get_type_hints
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.decorators import with_context_normalization
|
||||
|
||||
|
||||
class EnhancedBaseTool(BaseTool, ABC):
|
||||
"""Enhanced base class for MCP tools with automatic context normalization.
|
||||
|
||||
This base class automatically wraps the tool registration to include
|
||||
context normalization, ensuring that all tools handle external calls
|
||||
properly without requiring manual decoration or copy-pasted code.
|
||||
"""
|
||||
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this tool with automatic context normalization.
|
||||
|
||||
This method automatically applies context normalization to the tool
|
||||
handler, ensuring it works properly when called externally.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
"""
|
||||
# Get the tool method from the subclass
|
||||
tool_method = self._create_tool_handler()
|
||||
|
||||
# Apply context normalization decorator
|
||||
normalized_method = with_context_normalization(tool_method)
|
||||
|
||||
# Register with the server
|
||||
mcp_server.tool(name=self.name, description=self.description)(normalized_method)
|
||||
|
||||
@abstractmethod
|
||||
def _create_tool_handler(self) -> Any:
|
||||
"""Create the tool handler function.
|
||||
|
||||
Subclasses must implement this to return an async function
|
||||
that will be registered as the tool handler.
|
||||
|
||||
Returns:
|
||||
An async function that handles tool calls
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class AutoRegisterTool(BaseTool, ABC):
|
||||
"""Base class that automatically generates tool handlers from the call method.
|
||||
|
||||
This base class inspects the call method signature and automatically
|
||||
creates a properly typed tool handler with context normalization.
|
||||
"""
|
||||
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this tool with automatic handler generation.
|
||||
|
||||
This method inspects the call method signature and automatically
|
||||
creates a tool handler with the correct parameters and types.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
"""
|
||||
# Get the call method signature
|
||||
call_method = self.call
|
||||
sig = inspect.signature(call_method)
|
||||
|
||||
# Get type hints for proper typing
|
||||
hints = get_type_hints(call_method)
|
||||
|
||||
# Create a dynamic handler function
|
||||
tool_self = self
|
||||
|
||||
# Build the handler dynamically based on the call signature
|
||||
params = list(sig.parameters.items())
|
||||
|
||||
# Skip 'self' and 'ctx' parameters
|
||||
tool_params = [
|
||||
(name, param) for name, param in params if name not in ("self", "ctx")
|
||||
]
|
||||
|
||||
# Create the handler function dynamically
|
||||
async def handler(ctx: MCPContext, **kwargs: Any) -> Any:
|
||||
# Call the tool's call method with the context and parameters
|
||||
return await tool_self.call(ctx, **kwargs)
|
||||
|
||||
# Apply context normalization
|
||||
normalized_handler = with_context_normalization(handler)
|
||||
|
||||
# Register with the server
|
||||
mcp_server.tool(name=self.name, description=self.description)(
|
||||
normalized_handler
|
||||
)
|
||||
@@ -1,366 +0,0 @@
|
||||
"""FastMCP-compatible pagination implementation.
|
||||
|
||||
This module provides pagination utilities optimized for FastMCP with minimal latency.
|
||||
"""
|
||||
|
||||
import json
|
||||
import time
|
||||
import base64
|
||||
import hashlib
|
||||
from typing import Any, Dict, List, Union, Generic, TypeVar, Optional
|
||||
from dataclasses import field, dataclass
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@dataclass
|
||||
class CursorData:
|
||||
"""Cursor data structure for efficient pagination."""
|
||||
|
||||
# Primary cursor fields (indexed)
|
||||
last_id: Optional[str] = None
|
||||
last_timestamp: Optional[float] = None
|
||||
offset: int = 0
|
||||
|
||||
# Metadata for validation and optimization
|
||||
page_size: int = 100
|
||||
sort_field: str = "id"
|
||||
sort_order: str = "asc"
|
||||
|
||||
# Security and validation
|
||||
created_at: float = field(default_factory=time.time)
|
||||
expires_at: Optional[float] = None
|
||||
checksum: Optional[str] = None
|
||||
|
||||
def to_cursor(self) -> str:
|
||||
"""Convert to opaque cursor string."""
|
||||
data = {
|
||||
"id": self.last_id,
|
||||
"ts": self.last_timestamp,
|
||||
"o": self.offset,
|
||||
"ps": self.page_size,
|
||||
"sf": self.sort_field,
|
||||
"so": self.sort_order,
|
||||
"ca": self.created_at,
|
||||
}
|
||||
if self.expires_at:
|
||||
data["ea"] = self.expires_at
|
||||
|
||||
# Add checksum for integrity
|
||||
data_str = json.dumps(data, sort_keys=True, separators=(",", ":"))
|
||||
data["cs"] = hashlib.md5(data_str.encode()).hexdigest()[:8]
|
||||
|
||||
# Encode as base64
|
||||
final_str = json.dumps(data, separators=(",", ":"))
|
||||
return base64.urlsafe_b64encode(final_str.encode()).decode().rstrip("=")
|
||||
|
||||
@classmethod
|
||||
def from_cursor(cls, cursor: str) -> Optional["CursorData"]:
|
||||
"""Parse cursor string back to CursorData."""
|
||||
try:
|
||||
# Add padding if needed
|
||||
padding = 4 - (len(cursor) % 4)
|
||||
if padding != 4:
|
||||
cursor += "=" * padding
|
||||
|
||||
decoded = base64.urlsafe_b64decode(cursor.encode())
|
||||
data = json.loads(decoded)
|
||||
|
||||
# Validate checksum
|
||||
checksum = data.pop("cs", None)
|
||||
if checksum:
|
||||
data_str = json.dumps(
|
||||
{k: v for k, v in data.items() if k != "cs"},
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
)
|
||||
expected = hashlib.md5(data_str.encode()).hexdigest()[:8]
|
||||
if checksum != expected:
|
||||
return None
|
||||
|
||||
# Check expiration
|
||||
expires_at = data.get("ea")
|
||||
if expires_at and time.time() > expires_at:
|
||||
return None
|
||||
|
||||
return cls(
|
||||
last_id=data.get("id"),
|
||||
last_timestamp=data.get("ts"),
|
||||
offset=data.get("o", 0),
|
||||
page_size=data.get("ps", 100),
|
||||
sort_field=data.get("sf", "id"),
|
||||
sort_order=data.get("so", "asc"),
|
||||
created_at=data.get("ca", time.time()),
|
||||
expires_at=expires_at,
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
class FastMCPPaginator(Generic[T]):
|
||||
"""High-performance paginator for FastMCP responses."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
page_size: int = 100,
|
||||
max_page_size: int = 1000,
|
||||
cursor_ttl: int = 3600, # 1 hour
|
||||
enable_prefetch: bool = False,
|
||||
):
|
||||
"""Initialize the paginator.
|
||||
|
||||
Args:
|
||||
page_size: Default page size
|
||||
max_page_size: Maximum allowed page size
|
||||
cursor_ttl: Cursor time-to-live in seconds
|
||||
enable_prefetch: Enable prefetching for next page
|
||||
"""
|
||||
self.page_size = page_size
|
||||
self.max_page_size = max_page_size
|
||||
self.cursor_ttl = cursor_ttl
|
||||
self.enable_prefetch = enable_prefetch
|
||||
self._cache: Dict[str, Any] = {}
|
||||
|
||||
def paginate_list(
|
||||
self,
|
||||
items: List[T],
|
||||
cursor: Optional[str] = None,
|
||||
page_size: Optional[int] = None,
|
||||
sort_key: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Paginate a list with optimal performance.
|
||||
|
||||
Args:
|
||||
items: List to paginate
|
||||
cursor: Optional cursor from previous request
|
||||
page_size: Override default page size
|
||||
sort_key: Sort field for consistent ordering
|
||||
|
||||
Returns:
|
||||
Dict with items and optional nextCursor
|
||||
"""
|
||||
# Parse cursor or create new
|
||||
cursor_data = CursorData.from_cursor(cursor) if cursor else CursorData()
|
||||
|
||||
# Use provided page size or default
|
||||
actual_page_size = min(
|
||||
page_size or cursor_data.page_size or self.page_size, self.max_page_size
|
||||
)
|
||||
|
||||
# Get starting position
|
||||
start_idx = cursor_data.offset
|
||||
|
||||
# Validate bounds
|
||||
if start_idx >= len(items):
|
||||
return {"items": [], "hasMore": False}
|
||||
|
||||
# Slice the page
|
||||
end_idx = min(start_idx + actual_page_size, len(items))
|
||||
page_items = items[start_idx:end_idx]
|
||||
|
||||
# Build response
|
||||
response = {
|
||||
"items": page_items,
|
||||
"pageInfo": {
|
||||
"startIndex": start_idx,
|
||||
"endIndex": end_idx,
|
||||
"pageSize": len(page_items),
|
||||
"totalItems": len(items),
|
||||
},
|
||||
}
|
||||
|
||||
# Create next cursor if more items exist
|
||||
if end_idx < len(items):
|
||||
next_cursor_data = CursorData(
|
||||
offset=end_idx,
|
||||
page_size=actual_page_size,
|
||||
expires_at=time.time() + self.cursor_ttl if self.cursor_ttl else None,
|
||||
)
|
||||
response["nextCursor"] = next_cursor_data.to_cursor()
|
||||
response["hasMore"] = True
|
||||
else:
|
||||
response["hasMore"] = False
|
||||
|
||||
return response
|
||||
|
||||
def paginate_query(
|
||||
self,
|
||||
query_func,
|
||||
cursor: Optional[str] = None,
|
||||
page_size: Optional[int] = None,
|
||||
**query_params,
|
||||
) -> Dict[str, Any]:
|
||||
"""Paginate results from a query function.
|
||||
|
||||
This is optimized for database queries using indexed fields.
|
||||
|
||||
Args:
|
||||
query_func: Function that accepts (last_id, last_timestamp, limit, **params)
|
||||
cursor: Optional cursor
|
||||
page_size: Override page size
|
||||
**query_params: Additional query parameters
|
||||
|
||||
Returns:
|
||||
Paginated response
|
||||
"""
|
||||
# Parse cursor
|
||||
cursor_data = CursorData.from_cursor(cursor) if cursor else CursorData()
|
||||
|
||||
# Determine page size
|
||||
limit = min(
|
||||
page_size or cursor_data.page_size or self.page_size, self.max_page_size
|
||||
)
|
||||
|
||||
# Execute query with cursor position
|
||||
results = query_func(
|
||||
last_id=cursor_data.last_id,
|
||||
last_timestamp=cursor_data.last_timestamp,
|
||||
limit=limit + 1, # Fetch one extra to detect more
|
||||
**query_params,
|
||||
)
|
||||
|
||||
# Check if there are more results
|
||||
has_more = len(results) > limit
|
||||
if has_more:
|
||||
results = results[:limit] # Remove the extra item
|
||||
|
||||
# Build response
|
||||
response = {
|
||||
"items": results,
|
||||
"pageInfo": {"pageSize": len(results), "hasMore": has_more},
|
||||
}
|
||||
|
||||
# Create next cursor if needed
|
||||
if has_more and results:
|
||||
last_item = results[-1]
|
||||
next_cursor_data = CursorData(
|
||||
last_id=getattr(last_item, "id", None),
|
||||
last_timestamp=getattr(last_item, "timestamp", None),
|
||||
page_size=limit,
|
||||
sort_field=cursor_data.sort_field,
|
||||
sort_order=cursor_data.sort_order,
|
||||
expires_at=time.time() + self.cursor_ttl if self.cursor_ttl else None,
|
||||
)
|
||||
response["nextCursor"] = next_cursor_data.to_cursor()
|
||||
|
||||
return response
|
||||
|
||||
|
||||
class TokenAwarePaginator:
|
||||
"""Paginator that respects token limits for LLM responses."""
|
||||
|
||||
def __init__(self, max_tokens: int = 20000):
|
||||
"""Initialize token-aware paginator.
|
||||
|
||||
Args:
|
||||
max_tokens: Maximum tokens per response
|
||||
"""
|
||||
self.max_tokens = max_tokens
|
||||
self.paginator = FastMCPPaginator()
|
||||
|
||||
def paginate_by_tokens(
|
||||
self, items: List[Any], cursor: Optional[str] = None, estimate_func=None
|
||||
) -> Dict[str, Any]:
|
||||
"""Paginate items based on token count.
|
||||
|
||||
Args:
|
||||
items: Items to paginate
|
||||
cursor: Optional cursor
|
||||
estimate_func: Function to estimate tokens for an item
|
||||
|
||||
Returns:
|
||||
Paginated response
|
||||
"""
|
||||
from hanzo_mcp.tools.common.truncate import estimate_tokens
|
||||
|
||||
# Default token estimation
|
||||
if not estimate_func:
|
||||
estimate_func = lambda x: estimate_tokens(
|
||||
json.dumps(x) if not isinstance(x, str) else x
|
||||
)
|
||||
|
||||
# Parse cursor
|
||||
cursor_data = CursorData.from_cursor(cursor) if cursor else CursorData()
|
||||
start_idx = cursor_data.offset
|
||||
|
||||
# Build page respecting token limit
|
||||
page_items = []
|
||||
current_tokens = 100 # Base overhead
|
||||
current_idx = start_idx
|
||||
|
||||
while current_idx < len(items) and current_tokens < self.max_tokens:
|
||||
item = items[current_idx]
|
||||
item_tokens = estimate_func(item)
|
||||
|
||||
# Check if adding this item would exceed limit
|
||||
if current_tokens + item_tokens > self.max_tokens and page_items:
|
||||
break
|
||||
|
||||
page_items.append(item)
|
||||
current_tokens += item_tokens
|
||||
current_idx += 1
|
||||
|
||||
# Build response
|
||||
response = {
|
||||
"items": page_items,
|
||||
"pageInfo": {
|
||||
"itemCount": len(page_items),
|
||||
"estimatedTokens": current_tokens,
|
||||
"hasMore": current_idx < len(items),
|
||||
},
|
||||
}
|
||||
|
||||
# Add next cursor if needed
|
||||
if current_idx < len(items):
|
||||
next_cursor_data = CursorData(offset=current_idx)
|
||||
response["nextCursor"] = next_cursor_data.to_cursor()
|
||||
|
||||
return response
|
||||
|
||||
|
||||
# FastMCP integration helpers
|
||||
def create_paginated_response(
|
||||
items: Union[List[Any], Dict[str, Any], str],
|
||||
cursor: Optional[str] = None,
|
||||
page_size: int = 100,
|
||||
use_token_limit: bool = True,
|
||||
) -> Dict[str, Any]:
|
||||
"""Create a paginated response compatible with FastMCP.
|
||||
|
||||
Args:
|
||||
items: The items to paginate
|
||||
cursor: Optional cursor from request
|
||||
page_size: Items per page
|
||||
use_token_limit: Whether to use token-based pagination
|
||||
|
||||
Returns:
|
||||
FastMCP-compatible paginated response
|
||||
"""
|
||||
if use_token_limit:
|
||||
paginator = TokenAwarePaginator()
|
||||
|
||||
# Convert different types to list
|
||||
if isinstance(items, str):
|
||||
# Split string by lines for pagination
|
||||
items = items.split("\n")
|
||||
elif isinstance(items, dict):
|
||||
# Convert dict to list of key-value pairs
|
||||
items = [{"key": k, "value": v} for k, v in items.items()]
|
||||
|
||||
return paginator.paginate_by_tokens(items, cursor)
|
||||
else:
|
||||
paginator = FastMCPPaginator(page_size=page_size)
|
||||
|
||||
# Handle different input types
|
||||
if isinstance(items, list):
|
||||
return paginator.paginate_list(items, cursor, page_size)
|
||||
else:
|
||||
# Convert to list first
|
||||
if isinstance(items, str):
|
||||
items = items.split("\n")
|
||||
elif isinstance(items, dict):
|
||||
items = [{"key": k, "value": v} for k, v in items.items()]
|
||||
else:
|
||||
items = [items]
|
||||
|
||||
return paginator.paginate_list(items, cursor, page_size)
|
||||
@@ -1,257 +0,0 @@
|
||||
"""Forgiving edit helper for AI-friendly text matching."""
|
||||
|
||||
import re
|
||||
import difflib
|
||||
from typing import List, Tuple, Optional
|
||||
|
||||
|
||||
class ForgivingEditHelper:
|
||||
"""Helper class to make text editing more forgiving for AI usage.
|
||||
|
||||
This helper normalizes whitespace, handles partial matches, and provides
|
||||
suggestions when exact matches fail.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def normalize_whitespace(text: str) -> str:
|
||||
"""Normalize whitespace while preserving structure.
|
||||
|
||||
Args:
|
||||
text: Text to normalize
|
||||
|
||||
Returns:
|
||||
Text with normalized whitespace
|
||||
"""
|
||||
# Handle the input line by line
|
||||
lines = []
|
||||
for line in text.split("\n"):
|
||||
# Replace tabs with 4 spaces everywhere in the line
|
||||
line = line.replace("\t", " ")
|
||||
|
||||
# Split into indentation and content
|
||||
stripped = line.lstrip()
|
||||
indent = line[: len(line) - len(stripped)]
|
||||
|
||||
if stripped:
|
||||
# For content, normalize multiple spaces to single space
|
||||
content = re.sub(r" {2,}", " ", stripped)
|
||||
lines.append(indent + content)
|
||||
else:
|
||||
lines.append(indent)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
@staticmethod
|
||||
def find_fuzzy_match(
|
||||
haystack: str, needle: str, threshold: float = 0.85
|
||||
) -> Optional[Tuple[int, int, str]]:
|
||||
"""Find a fuzzy match for the needle in the haystack.
|
||||
|
||||
Args:
|
||||
haystack: Text to search in
|
||||
needle: Text to search for
|
||||
threshold: Similarity threshold (0-1)
|
||||
|
||||
Returns:
|
||||
Tuple of (start_pos, end_pos, matched_text) or None
|
||||
"""
|
||||
# First try exact match
|
||||
if needle in haystack:
|
||||
start = haystack.index(needle)
|
||||
return (start, start + len(needle), needle)
|
||||
|
||||
# Normalize for comparison
|
||||
norm_haystack = ForgivingEditHelper.normalize_whitespace(haystack)
|
||||
norm_needle = ForgivingEditHelper.normalize_whitespace(needle)
|
||||
|
||||
# Try normalized exact match
|
||||
if norm_needle in norm_haystack:
|
||||
# Find the match in normalized text
|
||||
norm_start = norm_haystack.index(norm_needle)
|
||||
|
||||
# Map back to original text
|
||||
# This is approximate but usually good enough
|
||||
lines_before = norm_haystack[:norm_start].count("\n")
|
||||
|
||||
# Find corresponding position in original
|
||||
original_lines = haystack.split("\n")
|
||||
norm_lines = norm_haystack.split("\n")
|
||||
|
||||
start_pos = sum(len(line) + 1 for line in original_lines[:lines_before])
|
||||
|
||||
# Find end position by counting lines in needle
|
||||
needle_lines = norm_needle.count("\n") + 1
|
||||
end_pos = sum(
|
||||
len(line) + 1 for line in original_lines[: lines_before + needle_lines]
|
||||
)
|
||||
|
||||
matched = "\n".join(
|
||||
original_lines[lines_before : lines_before + needle_lines]
|
||||
)
|
||||
return (start_pos, end_pos - 1, matched)
|
||||
|
||||
# Try fuzzy matching on lines
|
||||
haystack_lines = haystack.split("\n")
|
||||
needle_lines = needle.split("\n")
|
||||
|
||||
if len(needle_lines) == 1:
|
||||
# Single line - find best match
|
||||
needle_norm = ForgivingEditHelper.normalize_whitespace(needle)
|
||||
best_ratio = 0
|
||||
best_match = None
|
||||
|
||||
for i, line in enumerate(haystack_lines):
|
||||
line_norm = ForgivingEditHelper.normalize_whitespace(line)
|
||||
ratio = difflib.SequenceMatcher(None, line_norm, needle_norm).ratio()
|
||||
|
||||
if ratio > best_ratio and ratio >= threshold:
|
||||
best_ratio = ratio
|
||||
start_pos = sum(len(l) + 1 for l in haystack_lines[:i])
|
||||
best_match = (start_pos, start_pos + len(line), line)
|
||||
|
||||
return best_match
|
||||
|
||||
else:
|
||||
# Multi-line - find sequence match
|
||||
for i in range(len(haystack_lines) - len(needle_lines) + 1):
|
||||
candidate_lines = haystack_lines[i : i + len(needle_lines)]
|
||||
candidate = "\n".join(candidate_lines)
|
||||
candidate_norm = ForgivingEditHelper.normalize_whitespace(candidate)
|
||||
needle_norm = ForgivingEditHelper.normalize_whitespace(needle)
|
||||
|
||||
ratio = difflib.SequenceMatcher(
|
||||
None, candidate_norm, needle_norm
|
||||
).ratio()
|
||||
|
||||
if ratio >= threshold:
|
||||
start_pos = sum(len(l) + 1 for l in haystack_lines[:i])
|
||||
return (start_pos, start_pos + len(candidate), candidate)
|
||||
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def suggest_matches(
|
||||
haystack: str, needle: str, max_suggestions: int = 3
|
||||
) -> List[Tuple[float, str]]:
|
||||
"""Suggest possible matches when exact match fails.
|
||||
|
||||
Args:
|
||||
haystack: Text to search in
|
||||
needle: Text to search for
|
||||
max_suggestions: Maximum number of suggestions
|
||||
|
||||
Returns:
|
||||
List of (similarity_score, text) tuples
|
||||
"""
|
||||
suggestions = []
|
||||
|
||||
# Normalize needle
|
||||
needle_norm = ForgivingEditHelper.normalize_whitespace(needle)
|
||||
needle_lines = needle.split("\n")
|
||||
|
||||
if len(needle_lines) == 1:
|
||||
# Single line - compare with all lines
|
||||
for line in haystack.split("\n"):
|
||||
if line.strip(): # Skip empty lines
|
||||
line_norm = ForgivingEditHelper.normalize_whitespace(line)
|
||||
ratio = difflib.SequenceMatcher(
|
||||
None, line_norm, needle_norm
|
||||
).ratio()
|
||||
if ratio > 0.5: # Only reasonably similar lines
|
||||
suggestions.append((ratio, line))
|
||||
|
||||
else:
|
||||
# Multi-line - use sliding window
|
||||
haystack_lines = haystack.split("\n")
|
||||
window_size = len(needle_lines)
|
||||
|
||||
for i in range(len(haystack_lines) - window_size + 1):
|
||||
candidate_lines = haystack_lines[i : i + window_size]
|
||||
candidate = "\n".join(candidate_lines)
|
||||
candidate_norm = ForgivingEditHelper.normalize_whitespace(candidate)
|
||||
|
||||
ratio = difflib.SequenceMatcher(
|
||||
None, candidate_norm, needle_norm
|
||||
).ratio()
|
||||
if ratio > 0.5:
|
||||
suggestions.append((ratio, candidate))
|
||||
|
||||
# Sort by similarity and return top matches
|
||||
suggestions.sort(reverse=True, key=lambda x: x[0])
|
||||
return suggestions[:max_suggestions]
|
||||
|
||||
@staticmethod
|
||||
def create_edit_suggestion(
|
||||
file_content: str, old_string: str, new_string: str
|
||||
) -> dict:
|
||||
"""Create a helpful edit suggestion when match fails.
|
||||
|
||||
Args:
|
||||
file_content: Current file content
|
||||
old_string: String that couldn't be found
|
||||
new_string: Replacement string
|
||||
|
||||
Returns:
|
||||
Dict with error message and suggestions
|
||||
"""
|
||||
# Try fuzzy match
|
||||
fuzzy_match = ForgivingEditHelper.find_fuzzy_match(file_content, old_string)
|
||||
|
||||
if fuzzy_match:
|
||||
_, _, matched_text = fuzzy_match
|
||||
return {
|
||||
"error": "Exact match not found, but found similar text",
|
||||
"found": matched_text,
|
||||
"suggestion": "Use this as old_string instead",
|
||||
"confidence": "high",
|
||||
}
|
||||
|
||||
# Get suggestions
|
||||
suggestions = ForgivingEditHelper.suggest_matches(file_content, old_string)
|
||||
|
||||
if suggestions:
|
||||
return {
|
||||
"error": "Could not find exact or fuzzy match",
|
||||
"suggestions": [
|
||||
{"similarity": f"{score:.0%}", "text": text}
|
||||
for score, text in suggestions
|
||||
],
|
||||
"hint": "Try using one of these suggestions as old_string",
|
||||
}
|
||||
|
||||
# No good matches - provide general help
|
||||
return {
|
||||
"error": "Could not find any matches",
|
||||
"hints": [
|
||||
"Check for whitespace differences (tabs vs spaces)",
|
||||
"Ensure you're including complete lines",
|
||||
"Try a smaller, more unique portion of text",
|
||||
"Use the streaming_command tool to view the file with visible whitespace",
|
||||
],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def prepare_edit_string(text: str) -> str:
|
||||
"""Prepare a string for editing by handling common issues.
|
||||
|
||||
Args:
|
||||
text: Text to prepare
|
||||
|
||||
Returns:
|
||||
Cleaned text ready for editing
|
||||
"""
|
||||
# Remove any line number prefixes (common in AI copy-paste)
|
||||
lines = []
|
||||
for line in text.split("\n"):
|
||||
# Remove common line number patterns while preserving indentation
|
||||
# Match patterns like "1: ", "123: ", "1| ", "1- ", etc.
|
||||
# But preserve the original indentation after the line number
|
||||
match = re.match(r"^(\d+[:\|\-])\s(.*)", line)
|
||||
if match:
|
||||
# Keep only the content part (group 2) which includes any indentation
|
||||
lines.append(match.group(2))
|
||||
else:
|
||||
# No line number pattern found, keep the line as-is
|
||||
lines.append(line)
|
||||
|
||||
return "\n".join(lines)
|
||||
@@ -1,114 +0,0 @@
|
||||
"""Mode system for organizing development tools based on programmer personalities."""
|
||||
|
||||
import os
|
||||
from typing import Set, Dict, List, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
from hanzo_mcp.tools.common.personality import (
|
||||
ToolPersonality,
|
||||
personalities,
|
||||
ensure_agent_enabled,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Mode(ToolPersonality):
|
||||
"""Development mode combining tool preferences and environment settings."""
|
||||
|
||||
# Inherits all fields from ToolPersonality
|
||||
# Adds mode-specific functionality
|
||||
|
||||
@property
|
||||
def is_active(self) -> bool:
|
||||
"""Check if this mode is currently active."""
|
||||
return ModeRegistry.get_active() == self
|
||||
|
||||
|
||||
class ModeRegistry:
|
||||
"""Registry for development modes."""
|
||||
|
||||
_modes: Dict[str, Mode] = {}
|
||||
_active_mode: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def register(cls, mode: Mode) -> None:
|
||||
"""Register a development mode."""
|
||||
# Ensure agent is enabled if API keys present
|
||||
mode = ensure_agent_enabled(mode)
|
||||
cls._modes[mode.name] = mode
|
||||
|
||||
@classmethod
|
||||
def get(cls, name: str) -> Optional[Mode]:
|
||||
"""Get a mode by name."""
|
||||
return cls._modes.get(name)
|
||||
|
||||
@classmethod
|
||||
def list(cls) -> List[Mode]:
|
||||
"""List all registered modes."""
|
||||
return list(cls._modes.values())
|
||||
|
||||
@classmethod
|
||||
def set_active(cls, name: str) -> None:
|
||||
"""Set the active mode."""
|
||||
if name not in cls._modes:
|
||||
raise ValueError(f"Mode '{name}' not found")
|
||||
cls._active_mode = name
|
||||
|
||||
# Apply environment variables from the mode
|
||||
mode = cls._modes[name]
|
||||
if mode.environment:
|
||||
for key, value in mode.environment.items():
|
||||
os.environ[key] = value
|
||||
|
||||
@classmethod
|
||||
def get_active(cls) -> Optional[Mode]:
|
||||
"""Get the active mode."""
|
||||
if cls._active_mode:
|
||||
return cls._modes.get(cls._active_mode)
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def get_active_tools(cls) -> Set[str]:
|
||||
"""Get the set of tools from the active mode."""
|
||||
mode = cls.get_active()
|
||||
if mode:
|
||||
return set(mode.tools)
|
||||
return set()
|
||||
|
||||
|
||||
def register_default_modes():
|
||||
"""Register all default development modes."""
|
||||
# Convert personalities to modes
|
||||
for personality in personalities:
|
||||
mode = Mode(
|
||||
name=personality.name,
|
||||
programmer=personality.programmer,
|
||||
description=personality.description,
|
||||
tools=personality.tools,
|
||||
environment=personality.environment,
|
||||
philosophy=personality.philosophy,
|
||||
)
|
||||
ModeRegistry.register(mode)
|
||||
|
||||
|
||||
def get_mode_from_env() -> Optional[str]:
|
||||
"""Get mode name from environment variables."""
|
||||
# Check for HANZO_MODE, PERSONALITY, or MODE env vars
|
||||
return (
|
||||
os.environ.get("HANZO_MODE")
|
||||
or os.environ.get("PERSONALITY")
|
||||
or os.environ.get("MODE")
|
||||
)
|
||||
|
||||
|
||||
def activate_mode_from_env():
|
||||
"""Activate mode based on environment variables."""
|
||||
mode_name = get_mode_from_env()
|
||||
if mode_name:
|
||||
try:
|
||||
ModeRegistry.set_active(mode_name)
|
||||
return True
|
||||
except ValueError:
|
||||
# Mode not found, ignore
|
||||
pass
|
||||
return False
|
||||
@@ -1,108 +0,0 @@
|
||||
"""Tool mode loader for dynamic tool configuration."""
|
||||
|
||||
import os
|
||||
from typing import Dict, Optional
|
||||
|
||||
from hanzo_mcp.tools.common.mode import (
|
||||
ModeRegistry,
|
||||
activate_mode_from_env,
|
||||
register_default_modes,
|
||||
)
|
||||
|
||||
|
||||
class ModeLoader:
|
||||
"""Loads and manages tool modes for dynamic configuration."""
|
||||
|
||||
@staticmethod
|
||||
def initialize_modes() -> None:
|
||||
"""Initialize the mode system with defaults."""
|
||||
# Initialize modes
|
||||
register_default_modes()
|
||||
|
||||
# Check for mode from environment
|
||||
activate_mode_from_env()
|
||||
|
||||
# If no mode set, use default
|
||||
if not ModeRegistry.get_active():
|
||||
default_mode = os.environ.get("HANZO_DEFAULT_MODE", "hanzo")
|
||||
if ModeRegistry.get(default_mode):
|
||||
ModeRegistry.set_active(default_mode)
|
||||
|
||||
@staticmethod
|
||||
def get_enabled_tools_from_mode(
|
||||
base_enabled_tools: Optional[Dict[str, bool]] = None,
|
||||
force_mode: Optional[str] = None,
|
||||
) -> Dict[str, bool]:
|
||||
"""Get enabled tools configuration from active mode.
|
||||
|
||||
Args:
|
||||
base_enabled_tools: Base configuration to merge with
|
||||
force_mode: Force a specific mode (overrides active)
|
||||
|
||||
Returns:
|
||||
Dictionary of tool enable states
|
||||
"""
|
||||
# Initialize if needed
|
||||
if not ModeRegistry.list():
|
||||
ModeLoader.initialize_modes()
|
||||
|
||||
# Get mode to use
|
||||
tools_list = None
|
||||
|
||||
if force_mode:
|
||||
# Set and get mode
|
||||
if ModeRegistry.get(force_mode):
|
||||
ModeRegistry.set_active(force_mode)
|
||||
mode = ModeRegistry.get_active()
|
||||
tools_list = mode.tools if mode else None
|
||||
else:
|
||||
# Check active mode
|
||||
mode = ModeRegistry.get_active()
|
||||
if mode:
|
||||
tools_list = mode.tools
|
||||
|
||||
if not tools_list:
|
||||
# No active mode, return base config
|
||||
return base_enabled_tools or {}
|
||||
|
||||
# Start with base configuration
|
||||
result = base_enabled_tools.copy() if base_enabled_tools else {}
|
||||
|
||||
# Get all possible tools from registry
|
||||
from hanzo_mcp.config.tool_config import TOOL_REGISTRY
|
||||
|
||||
all_possible_tools = set(TOOL_REGISTRY.keys())
|
||||
|
||||
# Disable all tools first (clean slate for mode)
|
||||
for tool in all_possible_tools:
|
||||
result[tool] = False
|
||||
|
||||
# Enable tools from mode
|
||||
for tool in tools_list:
|
||||
result[tool] = True
|
||||
|
||||
# Always enable mode tool (meta)
|
||||
result["mode"] = True
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def get_environment_from_mode() -> Dict[str, str]:
|
||||
"""Get environment variables from active mode.
|
||||
|
||||
Returns:
|
||||
Dictionary of environment variables
|
||||
"""
|
||||
# Check mode
|
||||
mode = ModeRegistry.get_active()
|
||||
if mode and mode.environment:
|
||||
return mode.environment.copy()
|
||||
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
def apply_environment_from_mode() -> None:
|
||||
"""Apply environment variables from active mode."""
|
||||
env_vars = ModeLoader.get_environment_from_mode()
|
||||
for key, value in env_vars.items():
|
||||
os.environ[key] = value
|
||||
@@ -1,238 +0,0 @@
|
||||
"""Enhanced base class with automatic pagination support.
|
||||
|
||||
This module provides a base class that automatically handles pagination
|
||||
for all tool responses that exceed MCP token limits.
|
||||
"""
|
||||
|
||||
from abc import abstractmethod
|
||||
from typing import Any, Dict, Union
|
||||
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool, handle_connection_errors
|
||||
from hanzo_mcp.tools.common.pagination import CursorManager
|
||||
from hanzo_mcp.tools.common.paginated_response import paginate_if_needed
|
||||
|
||||
|
||||
class PaginatedBaseTool(BaseTool):
|
||||
"""Base class for tools with automatic pagination support.
|
||||
|
||||
This base class automatically handles pagination for responses that
|
||||
exceed MCP token limits, making all tools pagination-aware by default.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the paginated base tool."""
|
||||
super().__init__()
|
||||
self._supports_pagination = True
|
||||
|
||||
@abstractmethod
|
||||
async def execute(self, ctx: MCPContext, **params: Any) -> Any:
|
||||
"""Execute the tool logic and return raw results.
|
||||
|
||||
This method should be implemented by subclasses to perform the
|
||||
actual tool logic. The base class will handle pagination of
|
||||
the returned results automatically.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Raw tool results (will be paginated if needed)
|
||||
"""
|
||||
pass
|
||||
|
||||
@handle_connection_errors
|
||||
async def call(self, ctx: MCPContext, **params: Any) -> Union[str, Dict[str, Any]]:
|
||||
"""Execute the tool with automatic pagination support.
|
||||
|
||||
This method wraps the execute() method and automatically handles
|
||||
pagination if the response exceeds token limits.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters including optional 'cursor'
|
||||
|
||||
Returns:
|
||||
Tool result, potentially paginated
|
||||
"""
|
||||
# Extract cursor if provided
|
||||
cursor = params.pop("cursor", None)
|
||||
|
||||
# Validate cursor if provided
|
||||
if cursor and not CursorManager.parse_cursor(cursor):
|
||||
return {"error": "Invalid cursor provided", "code": -32602}
|
||||
|
||||
# Check if this is a continuation request
|
||||
if cursor:
|
||||
# For continuation, check if we have cached results
|
||||
cursor_data = CursorManager.parse_cursor(cursor)
|
||||
if (
|
||||
cursor_data
|
||||
and "tool" in cursor_data
|
||||
and cursor_data["tool"] != self.name
|
||||
):
|
||||
return {"error": "Cursor is for a different tool", "code": -32602}
|
||||
|
||||
# Execute the tool
|
||||
try:
|
||||
result = await self.execute(ctx, **params)
|
||||
except Exception as e:
|
||||
# Format errors consistently
|
||||
return {"error": str(e), "type": type(e).__name__}
|
||||
|
||||
# Handle pagination automatically
|
||||
if self._supports_pagination:
|
||||
paginated_result = paginate_if_needed(result, cursor)
|
||||
|
||||
# If pagination occurred, add tool info to help with continuation
|
||||
if isinstance(paginated_result, dict) and "nextCursor" in paginated_result:
|
||||
# Enhance the cursor with tool information
|
||||
if "nextCursor" in paginated_result:
|
||||
cursor_data = CursorManager.parse_cursor(
|
||||
paginated_result["nextCursor"]
|
||||
)
|
||||
if cursor_data:
|
||||
cursor_data["tool"] = self.name
|
||||
cursor_data["params"] = params # Store params for continuation
|
||||
paginated_result["nextCursor"] = CursorManager.create_cursor(
|
||||
cursor_data
|
||||
)
|
||||
|
||||
return paginated_result
|
||||
else:
|
||||
# Return raw result if pagination is disabled
|
||||
return result
|
||||
|
||||
def disable_pagination(self):
|
||||
"""Disable automatic pagination for this tool.
|
||||
|
||||
Some tools may want to handle their own pagination logic.
|
||||
"""
|
||||
self._supports_pagination = False
|
||||
|
||||
def enable_pagination(self):
|
||||
"""Re-enable automatic pagination for this tool."""
|
||||
self._supports_pagination = True
|
||||
|
||||
|
||||
class PaginatedFileSystemTool(PaginatedBaseTool):
|
||||
"""Base class for filesystem tools with pagination support."""
|
||||
|
||||
def __init__(self, permission_manager):
|
||||
"""Initialize filesystem tool with pagination.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
"""
|
||||
super().__init__()
|
||||
self.permission_manager = permission_manager
|
||||
|
||||
def is_path_allowed(self, path: str) -> bool:
|
||||
"""Check if a path is allowed according to permission settings.
|
||||
|
||||
Args:
|
||||
path: Path to check
|
||||
|
||||
Returns:
|
||||
True if the path is allowed, False otherwise
|
||||
"""
|
||||
return self.permission_manager.is_path_allowed(path)
|
||||
|
||||
|
||||
def migrate_tool_to_paginated(tool_class):
|
||||
"""Decorator to migrate existing tools to use pagination.
|
||||
|
||||
This decorator can be applied to existing tool classes to add
|
||||
automatic pagination support without modifying their code.
|
||||
|
||||
Usage:
|
||||
@migrate_tool_to_paginated
|
||||
class MyTool(BaseTool):
|
||||
...
|
||||
"""
|
||||
|
||||
class PaginatedWrapper(PaginatedBaseTool):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__()
|
||||
self._wrapped_tool = tool_class(*args, **kwargs)
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
return self._wrapped_tool.name
|
||||
|
||||
@property
|
||||
def description(self):
|
||||
# Add pagination info to description
|
||||
desc = self._wrapped_tool.description
|
||||
if not "pagination" in desc.lower():
|
||||
desc += "\n\nThis tool supports automatic pagination. If the response is too large, it will be split across multiple requests. Use the returned cursor to continue."
|
||||
return desc
|
||||
|
||||
async def execute(self, ctx: MCPContext, **params: Any) -> Any:
|
||||
# Call the wrapped tool's call method
|
||||
return await self._wrapped_tool.call(ctx, **params)
|
||||
|
||||
def register(self, mcp_server):
|
||||
# Need to create a new registration that includes cursor parameter
|
||||
tool_self = self
|
||||
|
||||
# Get the original registration function
|
||||
original_register = self._wrapped_tool.register
|
||||
|
||||
# Create a new registration that adds cursor support
|
||||
def register_with_pagination(server):
|
||||
# First register the original tool
|
||||
original_register(server)
|
||||
|
||||
# Then override with pagination support
|
||||
import inspect
|
||||
|
||||
# Get the registered function
|
||||
tool_func = None
|
||||
for name, func in server._tools.items():
|
||||
if name == self.name:
|
||||
tool_func = func
|
||||
break
|
||||
|
||||
if tool_func:
|
||||
# Get original signature
|
||||
sig = inspect.signature(tool_func)
|
||||
params = list(sig.parameters.values())
|
||||
|
||||
# Add cursor parameter if not present
|
||||
has_cursor = any(p.name == "cursor" for p in params)
|
||||
if not has_cursor:
|
||||
import inspect
|
||||
from typing import Optional
|
||||
|
||||
# Create new parameter with cursor
|
||||
cursor_param = inspect.Parameter(
|
||||
"cursor",
|
||||
inspect.Parameter.KEYWORD_ONLY,
|
||||
default=None,
|
||||
annotation=Optional[str],
|
||||
)
|
||||
|
||||
# Insert before ctx parameter
|
||||
new_params = []
|
||||
for p in params:
|
||||
if p.name == "ctx":
|
||||
new_params.append(cursor_param)
|
||||
new_params.append(p)
|
||||
|
||||
# Create wrapper function
|
||||
async def paginated_wrapper(**kwargs):
|
||||
return await tool_self.call(kwargs.get("ctx"), **kwargs)
|
||||
|
||||
# Update registration
|
||||
server._tools[self.name] = paginated_wrapper
|
||||
|
||||
register_with_pagination(mcp_server)
|
||||
|
||||
# Set the class name
|
||||
PaginatedWrapper.__name__ = f"Paginated{tool_class.__name__}"
|
||||
PaginatedWrapper.__qualname__ = f"Paginated{tool_class.__qualname__}"
|
||||
|
||||
return PaginatedWrapper
|
||||
@@ -1,300 +0,0 @@
|
||||
"""Automatic pagination response wrapper for MCP tools.
|
||||
|
||||
This module provides utilities to automatically paginate tool responses
|
||||
when they exceed token limits.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, Union, Optional
|
||||
|
||||
from hanzo_mcp.tools.common.truncate import estimate_tokens
|
||||
from hanzo_mcp.tools.common.pagination import CursorManager
|
||||
|
||||
|
||||
class AutoPaginatedResponse:
|
||||
"""Automatically paginate responses that exceed token limits."""
|
||||
|
||||
# MCP token limit with safety buffer
|
||||
MAX_TOKENS = 20000 # Leave 5k buffer from 25k limit
|
||||
|
||||
@staticmethod
|
||||
def create_response(
|
||||
content: Union[str, Dict[str, Any], List[Any]],
|
||||
cursor: Optional[str] = None,
|
||||
max_tokens: int = MAX_TOKENS,
|
||||
) -> Dict[str, Any]:
|
||||
"""Create a response that automatically paginates if too large.
|
||||
|
||||
Args:
|
||||
content: The content to return
|
||||
cursor: Optional cursor from request
|
||||
max_tokens: Maximum tokens allowed in response
|
||||
|
||||
Returns:
|
||||
Dict with content and optional nextCursor
|
||||
"""
|
||||
# Handle different content types
|
||||
if isinstance(content, str):
|
||||
return AutoPaginatedResponse._handle_string_response(
|
||||
content, cursor, max_tokens
|
||||
)
|
||||
elif isinstance(content, list):
|
||||
return AutoPaginatedResponse._handle_list_response(
|
||||
content, cursor, max_tokens
|
||||
)
|
||||
elif isinstance(content, dict):
|
||||
# If dict already has pagination info, return as-is
|
||||
if "nextCursor" in content or "cursor" in content:
|
||||
return content
|
||||
# Otherwise treat as single item
|
||||
return AutoPaginatedResponse._handle_dict_response(
|
||||
content, cursor, max_tokens
|
||||
)
|
||||
else:
|
||||
# Convert to string for other types
|
||||
return AutoPaginatedResponse._handle_string_response(
|
||||
str(content), cursor, max_tokens
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _handle_string_response(
|
||||
content: str, cursor: Optional[str], max_tokens: int
|
||||
) -> Dict[str, Any]:
|
||||
"""Handle pagination for string responses."""
|
||||
# Parse cursor to get offset
|
||||
offset = 0
|
||||
if cursor:
|
||||
cursor_data = CursorManager.parse_cursor(cursor)
|
||||
if cursor_data and "offset" in cursor_data:
|
||||
offset = cursor_data["offset"]
|
||||
|
||||
# For strings, paginate by lines
|
||||
lines = content.split("\n")
|
||||
|
||||
if offset >= len(lines):
|
||||
return {"content": "", "message": "No more content"}
|
||||
|
||||
# Build response line by line, checking tokens
|
||||
result_lines = []
|
||||
current_tokens = 0
|
||||
line_index = offset
|
||||
|
||||
# Add header if this is a continuation
|
||||
if offset > 0:
|
||||
header = f"[Continued from line {offset + 1}]\n"
|
||||
current_tokens = estimate_tokens(header)
|
||||
result_lines.append(header)
|
||||
|
||||
while line_index < len(lines):
|
||||
line = lines[line_index]
|
||||
line_tokens = estimate_tokens(line + "\n")
|
||||
|
||||
# Check if adding this line would exceed limit
|
||||
if current_tokens + line_tokens > max_tokens:
|
||||
# Need to paginate
|
||||
if not result_lines:
|
||||
# Single line too long, truncate it
|
||||
truncated_line = line[:1000] + "... [line truncated]"
|
||||
result_lines.append(truncated_line)
|
||||
line_index += 1
|
||||
break
|
||||
|
||||
result_lines.append(line)
|
||||
current_tokens += line_tokens
|
||||
line_index += 1
|
||||
|
||||
# Build response
|
||||
response = {"content": "\n".join(result_lines)}
|
||||
|
||||
# Add pagination info
|
||||
if line_index < len(lines):
|
||||
response["nextCursor"] = CursorManager.create_offset_cursor(line_index)
|
||||
response["pagination_info"] = {
|
||||
"current_lines": f"{offset + 1}-{line_index}",
|
||||
"total_lines": len(lines),
|
||||
"has_more": True,
|
||||
}
|
||||
else:
|
||||
response["pagination_info"] = {
|
||||
"current_lines": f"{offset + 1}-{len(lines)}",
|
||||
"total_lines": len(lines),
|
||||
"has_more": False,
|
||||
}
|
||||
|
||||
return response
|
||||
|
||||
@staticmethod
|
||||
def _handle_list_response(
|
||||
items: List[Any], cursor: Optional[str], max_tokens: int
|
||||
) -> Dict[str, Any]:
|
||||
"""Handle pagination for list responses."""
|
||||
# Parse cursor to get offset
|
||||
offset = 0
|
||||
if cursor:
|
||||
cursor_data = CursorManager.parse_cursor(cursor)
|
||||
if cursor_data and "offset" in cursor_data:
|
||||
offset = cursor_data["offset"]
|
||||
|
||||
if offset >= len(items):
|
||||
return {"items": [], "message": "No more items"}
|
||||
|
||||
# Build response item by item, checking tokens
|
||||
result_items = []
|
||||
current_tokens = 100 # Base overhead
|
||||
item_index = offset
|
||||
|
||||
# Add header if continuation
|
||||
header_obj = {}
|
||||
if offset > 0:
|
||||
header_obj["continuation_from"] = offset
|
||||
current_tokens += 50
|
||||
|
||||
while item_index < len(items):
|
||||
item = items[item_index]
|
||||
|
||||
# Estimate tokens for this item
|
||||
item_str = json.dumps(item) if not isinstance(item, str) else item
|
||||
item_tokens = estimate_tokens(item_str)
|
||||
|
||||
# Check if adding this item would exceed limit
|
||||
if current_tokens + item_tokens > max_tokens:
|
||||
if not result_items:
|
||||
# Single item too large, truncate it
|
||||
if isinstance(item, str):
|
||||
truncated = item[:5000] + "... [truncated]"
|
||||
result_items.append(truncated)
|
||||
else:
|
||||
result_items.append(
|
||||
{"error": "Item too large", "index": item_index}
|
||||
)
|
||||
item_index += 1
|
||||
break
|
||||
|
||||
result_items.append(item)
|
||||
current_tokens += item_tokens
|
||||
item_index += 1
|
||||
|
||||
# Build response
|
||||
response = {"items": result_items}
|
||||
|
||||
if header_obj:
|
||||
response.update(header_obj)
|
||||
|
||||
# Add pagination info
|
||||
if item_index < len(items):
|
||||
response["nextCursor"] = CursorManager.create_offset_cursor(item_index)
|
||||
response["pagination_info"] = {
|
||||
"returned_items": len(result_items),
|
||||
"total_items": len(items),
|
||||
"has_more": True,
|
||||
"next_index": item_index,
|
||||
}
|
||||
else:
|
||||
response["pagination_info"] = {
|
||||
"returned_items": len(result_items),
|
||||
"total_items": len(items),
|
||||
"has_more": False,
|
||||
}
|
||||
|
||||
return response
|
||||
|
||||
@staticmethod
|
||||
def _handle_dict_response(
|
||||
content: Dict[str, Any], cursor: Optional[str], max_tokens: int
|
||||
) -> Dict[str, Any]:
|
||||
"""Handle pagination for dict responses."""
|
||||
# For dicts, check if it's too large as-is
|
||||
content_str = json.dumps(content, indent=2)
|
||||
content_tokens = estimate_tokens(content_str)
|
||||
|
||||
if content_tokens <= max_tokens:
|
||||
# Fits within limit
|
||||
return content
|
||||
|
||||
# Too large - need to paginate
|
||||
# Strategy: Convert to key-value pairs and paginate
|
||||
items = list(content.items())
|
||||
offset = 0
|
||||
|
||||
if cursor:
|
||||
cursor_data = CursorManager.parse_cursor(cursor)
|
||||
if cursor_data and "offset" in cursor_data:
|
||||
offset = cursor_data["offset"]
|
||||
|
||||
if offset >= len(items):
|
||||
return {"content": {}, "message": "No more content"}
|
||||
|
||||
# Build paginated dict
|
||||
result = {}
|
||||
current_tokens = 100 # Base overhead
|
||||
|
||||
for i in range(offset, len(items)):
|
||||
key, value = items[i]
|
||||
|
||||
# Estimate tokens for this entry
|
||||
entry_str = json.dumps({key: value})
|
||||
entry_tokens = estimate_tokens(entry_str)
|
||||
|
||||
if current_tokens + entry_tokens > max_tokens:
|
||||
if not result:
|
||||
# Single entry too large
|
||||
result[key] = "[Value too large - use specific key access]"
|
||||
break
|
||||
|
||||
result[key] = value
|
||||
current_tokens += entry_tokens
|
||||
|
||||
# Wrap in response
|
||||
response = {"content": result}
|
||||
|
||||
# Add pagination info
|
||||
processed = offset + len(result)
|
||||
if processed < len(items):
|
||||
response["nextCursor"] = CursorManager.create_offset_cursor(processed)
|
||||
response["pagination_info"] = {
|
||||
"keys_returned": len(result),
|
||||
"total_keys": len(items),
|
||||
"has_more": True,
|
||||
}
|
||||
else:
|
||||
response["pagination_info"] = {
|
||||
"keys_returned": len(result),
|
||||
"total_keys": len(items),
|
||||
"has_more": False,
|
||||
}
|
||||
|
||||
return response
|
||||
|
||||
|
||||
def paginate_if_needed(
|
||||
response: Any, cursor: Optional[str] = None, force_pagination: bool = False
|
||||
) -> Union[str, Dict[str, Any]]:
|
||||
"""Wrap a response with automatic pagination if needed.
|
||||
|
||||
Args:
|
||||
response: The response to potentially paginate
|
||||
cursor: Optional cursor from request
|
||||
force_pagination: Force pagination even for small responses
|
||||
|
||||
Returns:
|
||||
Original response if small enough, otherwise paginated dict
|
||||
"""
|
||||
# Quick check - if response is already paginated, return as-is
|
||||
if isinstance(response, dict) and (
|
||||
"nextCursor" in response or "pagination_info" in response
|
||||
):
|
||||
return response
|
||||
|
||||
# For small responses, don't paginate unless forced
|
||||
if not force_pagination:
|
||||
try:
|
||||
response_str = (
|
||||
json.dumps(response) if not isinstance(response, str) else response
|
||||
)
|
||||
if len(response_str) < 10000: # Quick heuristic
|
||||
return response
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Create paginated response
|
||||
return AutoPaginatedResponse.create_response(response, cursor)
|
||||
@@ -1,223 +0,0 @@
|
||||
"""Pagination utilities for MCP tools.
|
||||
|
||||
This module provides utilities for implementing cursor-based pagination
|
||||
according to the MCP pagination protocol.
|
||||
"""
|
||||
|
||||
import json
|
||||
import base64
|
||||
from typing import Any, Dict, List, Generic, TypeVar, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@dataclass
|
||||
class PaginationParams:
|
||||
"""Parameters for pagination."""
|
||||
|
||||
cursor: Optional[str] = None
|
||||
page_size: int = 100 # Default page size
|
||||
|
||||
|
||||
@dataclass
|
||||
class PaginatedResponse(Generic[T]):
|
||||
"""A paginated response containing items and optional next cursor."""
|
||||
|
||||
items: List[T]
|
||||
next_cursor: Optional[str] = None
|
||||
|
||||
def to_dict(self, items_key: str = "items") -> Dict[str, Any]:
|
||||
"""Convert to dictionary format for MCP response.
|
||||
|
||||
Args:
|
||||
items_key: The key to use for items in the response
|
||||
|
||||
Returns:
|
||||
Dictionary with items and optional nextCursor
|
||||
"""
|
||||
result = {items_key: self.items}
|
||||
if self.next_cursor:
|
||||
result["nextCursor"] = self.next_cursor
|
||||
return result
|
||||
|
||||
|
||||
class CursorManager:
|
||||
"""Manages cursor creation and parsing for pagination."""
|
||||
|
||||
@staticmethod
|
||||
def create_cursor(data: Dict[str, Any]) -> str:
|
||||
"""Create an opaque cursor from data.
|
||||
|
||||
Args:
|
||||
data: Data to encode in the cursor
|
||||
|
||||
Returns:
|
||||
Base64-encoded cursor string
|
||||
"""
|
||||
json_data = json.dumps(data, separators=(",", ":"))
|
||||
return base64.b64encode(json_data.encode()).decode()
|
||||
|
||||
@staticmethod
|
||||
def parse_cursor(cursor: str) -> Optional[Dict[str, Any]]:
|
||||
"""Parse a cursor string back to data.
|
||||
|
||||
Args:
|
||||
cursor: Base64-encoded cursor string
|
||||
|
||||
Returns:
|
||||
Decoded data or None if invalid
|
||||
"""
|
||||
try:
|
||||
decoded = base64.b64decode(cursor.encode()).decode()
|
||||
return json.loads(decoded)
|
||||
except (ValueError, json.JSONDecodeError):
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def create_offset_cursor(offset: int) -> str:
|
||||
"""Create a cursor for offset-based pagination.
|
||||
|
||||
Args:
|
||||
offset: The offset for the next page
|
||||
|
||||
Returns:
|
||||
Cursor string
|
||||
"""
|
||||
return CursorManager.create_cursor({"offset": offset})
|
||||
|
||||
@staticmethod
|
||||
def parse_offset_cursor(cursor: Optional[str]) -> int:
|
||||
"""Parse an offset cursor.
|
||||
|
||||
Args:
|
||||
cursor: Cursor string or None
|
||||
|
||||
Returns:
|
||||
Offset value (0 if cursor is None or invalid)
|
||||
"""
|
||||
if not cursor:
|
||||
return 0
|
||||
|
||||
data = CursorManager.parse_cursor(cursor)
|
||||
if not data or "offset" not in data:
|
||||
return 0
|
||||
|
||||
return int(data["offset"])
|
||||
|
||||
|
||||
class Paginator(Generic[T]):
|
||||
"""Generic paginator for any list of items."""
|
||||
|
||||
def __init__(self, items: List[T], page_size: int = 100):
|
||||
"""Initialize paginator.
|
||||
|
||||
Args:
|
||||
items: List of items to paginate
|
||||
page_size: Number of items per page
|
||||
"""
|
||||
self.items = items
|
||||
self.page_size = page_size
|
||||
|
||||
def get_page(self, cursor: Optional[str] = None) -> PaginatedResponse[T]:
|
||||
"""Get a page of results.
|
||||
|
||||
Args:
|
||||
cursor: Optional cursor for the page
|
||||
|
||||
Returns:
|
||||
Paginated response with items and next cursor
|
||||
"""
|
||||
offset = CursorManager.parse_offset_cursor(cursor)
|
||||
|
||||
# Get the page of items
|
||||
start = offset
|
||||
end = min(start + self.page_size, len(self.items))
|
||||
page_items = self.items[start:end]
|
||||
|
||||
# Create next cursor if there are more items
|
||||
next_cursor = None
|
||||
if end < len(self.items):
|
||||
next_cursor = CursorManager.create_offset_cursor(end)
|
||||
|
||||
return PaginatedResponse(items=page_items, next_cursor=next_cursor)
|
||||
|
||||
|
||||
class StreamPaginator(Generic[T]):
|
||||
"""Paginator for streaming/generator-based results."""
|
||||
|
||||
def __init__(self, page_size: int = 100):
|
||||
"""Initialize stream paginator.
|
||||
|
||||
Args:
|
||||
page_size: Number of items per page
|
||||
"""
|
||||
self.page_size = page_size
|
||||
|
||||
def paginate_stream(
|
||||
self, stream_generator, cursor: Optional[str] = None
|
||||
) -> PaginatedResponse[T]:
|
||||
"""Paginate results from a stream/generator.
|
||||
|
||||
Args:
|
||||
stream_generator: Generator function that yields items
|
||||
cursor: Optional cursor for resuming
|
||||
|
||||
Returns:
|
||||
Paginated response
|
||||
"""
|
||||
items = []
|
||||
skip_count = 0
|
||||
|
||||
# Parse cursor to get skip count
|
||||
if cursor:
|
||||
cursor_data = CursorManager.parse_cursor(cursor)
|
||||
if cursor_data and "skip" in cursor_data:
|
||||
skip_count = cursor_data["skip"]
|
||||
|
||||
# Skip items based on cursor
|
||||
item_count = 0
|
||||
for item in stream_generator():
|
||||
if item_count < skip_count:
|
||||
item_count += 1
|
||||
continue
|
||||
|
||||
items.append(item)
|
||||
if len(items) >= self.page_size:
|
||||
# We have a full page, create cursor for next page
|
||||
next_cursor = CursorManager.create_cursor(
|
||||
{"skip": skip_count + len(items)}
|
||||
)
|
||||
return PaginatedResponse(items=items, next_cursor=next_cursor)
|
||||
|
||||
# No more items
|
||||
return PaginatedResponse(items=items, next_cursor=None)
|
||||
|
||||
|
||||
def paginate_list(
|
||||
items: List[T], cursor: Optional[str] = None, page_size: int = 100
|
||||
) -> PaginatedResponse[T]:
|
||||
"""Convenience function to paginate a list.
|
||||
|
||||
Args:
|
||||
items: List to paginate
|
||||
cursor: Optional cursor
|
||||
page_size: Items per page
|
||||
|
||||
Returns:
|
||||
Paginated response
|
||||
"""
|
||||
paginator = Paginator(items, page_size)
|
||||
return paginator.get_page(cursor)
|
||||
|
||||
|
||||
def validate_cursor(cursor: str) -> bool:
|
||||
"""Validate that a cursor is properly formatted.
|
||||
|
||||
Args:
|
||||
cursor: Cursor string to validate
|
||||
|
||||
Returns:
|
||||
True if valid, False otherwise
|
||||
"""
|
||||
return CursorManager.parse_cursor(cursor) is not None
|
||||
@@ -1,265 +0,0 @@
|
||||
"""Permission system for the Hanzo AI server."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import tempfile
|
||||
from typing import Any, TypeVar, final
|
||||
from pathlib import Path
|
||||
from collections.abc import Callable, Awaitable
|
||||
|
||||
# Define type variables for better type annotations
|
||||
T = TypeVar("T")
|
||||
P = TypeVar("P")
|
||||
|
||||
|
||||
@final
|
||||
class PermissionManager:
|
||||
"""Manages permissions for file and command operations."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialize the permission manager."""
|
||||
# Allowed paths
|
||||
self.allowed_paths: set[Path] = set()
|
||||
|
||||
# Allowed paths based on platform
|
||||
if sys.platform == "win32": # Windows
|
||||
self.allowed_paths.add(Path(tempfile.gettempdir()).resolve())
|
||||
else: # Unix/Linux/Mac
|
||||
self.allowed_paths.add(Path("/tmp").resolve())
|
||||
self.allowed_paths.add(Path("/var").resolve())
|
||||
|
||||
# Also allow user's home directory work folders
|
||||
home = Path.home()
|
||||
if home.exists():
|
||||
# Add common development directories
|
||||
work_dir = home / "work"
|
||||
if work_dir.exists():
|
||||
self.allowed_paths.add(work_dir.resolve())
|
||||
|
||||
# Excluded paths
|
||||
self.excluded_paths: set[Path] = set()
|
||||
self.excluded_patterns: list[str] = []
|
||||
|
||||
# Default excluded patterns
|
||||
self._add_default_exclusions()
|
||||
|
||||
def _add_default_exclusions(self) -> None:
|
||||
"""Add default exclusions for sensitive files and directories."""
|
||||
# Sensitive directories
|
||||
sensitive_dirs: list[str] = [
|
||||
".ssh",
|
||||
".gnupg",
|
||||
"node_modules",
|
||||
"__pycache__",
|
||||
".venv",
|
||||
"venv",
|
||||
"env",
|
||||
".idea",
|
||||
".DS_Store",
|
||||
]
|
||||
self.excluded_patterns.extend(sensitive_dirs)
|
||||
|
||||
# Sensitive file patterns
|
||||
sensitive_patterns: list[str] = [
|
||||
".env",
|
||||
"*.key",
|
||||
"*.pem",
|
||||
"*.crt",
|
||||
"*password*",
|
||||
"*secret*",
|
||||
"*.sqlite",
|
||||
"*.db",
|
||||
"*.sqlite3",
|
||||
"*.log",
|
||||
]
|
||||
self.excluded_patterns.extend(sensitive_patterns)
|
||||
|
||||
def add_allowed_path(self, path: str) -> None:
|
||||
"""Add a path to the allowed paths.
|
||||
|
||||
Args:
|
||||
path: The path to allow
|
||||
"""
|
||||
resolved_path: Path = Path(path).resolve()
|
||||
self.allowed_paths.add(resolved_path)
|
||||
|
||||
def remove_allowed_path(self, path: str) -> None:
|
||||
"""Remove a path from the allowed paths.
|
||||
|
||||
Args:
|
||||
path: The path to remove
|
||||
"""
|
||||
resolved_path: Path = Path(path).resolve()
|
||||
if resolved_path in self.allowed_paths:
|
||||
self.allowed_paths.remove(resolved_path)
|
||||
|
||||
def exclude_path(self, path: str) -> None:
|
||||
"""Exclude a path from allowed operations.
|
||||
|
||||
Args:
|
||||
path: The path to exclude
|
||||
"""
|
||||
resolved_path: Path = Path(path).resolve()
|
||||
self.excluded_paths.add(resolved_path)
|
||||
|
||||
def add_exclusion_pattern(self, pattern: str) -> None:
|
||||
"""Add an exclusion pattern.
|
||||
|
||||
Args:
|
||||
pattern: The pattern to exclude
|
||||
"""
|
||||
self.excluded_patterns.append(pattern)
|
||||
|
||||
def is_path_allowed(self, path: str) -> bool:
|
||||
"""Check if a path is allowed.
|
||||
|
||||
Args:
|
||||
path: The path to check
|
||||
|
||||
Returns:
|
||||
True if the path is allowed, False otherwise
|
||||
"""
|
||||
resolved_path: Path = Path(path).resolve()
|
||||
|
||||
# Check exclusions first
|
||||
if self._is_path_excluded(resolved_path):
|
||||
return False
|
||||
|
||||
# Check if the path is within any allowed path
|
||||
for allowed_path in self.allowed_paths:
|
||||
try:
|
||||
resolved_path.relative_to(allowed_path)
|
||||
return True
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
return False
|
||||
|
||||
def _is_path_excluded(self, path: Path) -> bool:
|
||||
"""Check if a path is excluded.
|
||||
|
||||
Args:
|
||||
path: The path to check
|
||||
|
||||
Returns:
|
||||
True if the path is excluded, False otherwise
|
||||
"""
|
||||
|
||||
# Check exact excluded paths
|
||||
if path in self.excluded_paths:
|
||||
return True
|
||||
|
||||
# Check excluded patterns
|
||||
path_str: str = str(path)
|
||||
|
||||
# Get path parts to check for exact directory/file name matches
|
||||
path_parts = path_str.split(os.sep)
|
||||
|
||||
for pattern in self.excluded_patterns:
|
||||
# Handle wildcard patterns (e.g., "*.log")
|
||||
if pattern.startswith("*"):
|
||||
if path_str.endswith(pattern[1:]):
|
||||
return True
|
||||
else:
|
||||
# For non-wildcard patterns, check if any path component matches exactly
|
||||
if pattern in path_parts:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def to_json(self) -> str:
|
||||
"""Convert the permission manager to a JSON string.
|
||||
|
||||
Returns:
|
||||
A JSON string representation of the permission manager
|
||||
"""
|
||||
data: dict[str, Any] = {
|
||||
"allowed_paths": [str(p) for p in self.allowed_paths],
|
||||
"excluded_paths": [str(p) for p in self.excluded_paths],
|
||||
"excluded_patterns": self.excluded_patterns,
|
||||
}
|
||||
|
||||
return json.dumps(data)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, json_str: str) -> "PermissionManager":
|
||||
"""Create a permission manager from a JSON string.
|
||||
|
||||
Args:
|
||||
json_str: The JSON string
|
||||
|
||||
Returns:
|
||||
A new PermissionManager instance
|
||||
"""
|
||||
data: dict[str, Any] = json.loads(json_str)
|
||||
|
||||
manager = cls()
|
||||
|
||||
for path in data.get("allowed_paths", []):
|
||||
manager.add_allowed_path(path)
|
||||
|
||||
for path in data.get("excluded_paths", []):
|
||||
manager.exclude_path(path)
|
||||
|
||||
manager.excluded_patterns = data.get("excluded_patterns", [])
|
||||
|
||||
return manager
|
||||
|
||||
|
||||
class PermissibleOperation:
|
||||
"""A decorator for operations that require permission."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
operation: str,
|
||||
get_path_fn: Callable[[list[Any], dict[str, Any]], str] | None = None,
|
||||
) -> None:
|
||||
"""Initialize the permissible operation.
|
||||
|
||||
Args:
|
||||
permission_manager: The permission manager
|
||||
operation: The operation type (read, write, execute, etc.)
|
||||
get_path_fn: Optional function to extract the path from args and kwargs
|
||||
"""
|
||||
self.permission_manager: PermissionManager = permission_manager
|
||||
self.operation: str = operation
|
||||
self.get_path_fn: Callable[[list[Any], dict[str, Any]], str] | None = (
|
||||
get_path_fn
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self, func: Callable[..., Awaitable[T]]
|
||||
) -> Callable[..., Awaitable[T]]:
|
||||
"""Decorate the function.
|
||||
|
||||
Args:
|
||||
func: The function to decorate
|
||||
|
||||
Returns:
|
||||
The decorated function
|
||||
"""
|
||||
|
||||
async def wrapper(*args: Any, **kwargs: Any) -> T:
|
||||
# Extract the path
|
||||
if self.get_path_fn:
|
||||
# Pass args as a list and kwargs as a dict to the path function
|
||||
path = self.get_path_fn(list(args), kwargs)
|
||||
else:
|
||||
# Default to first argument
|
||||
path = args[0] if args else next(iter(kwargs.values()), None)
|
||||
|
||||
if not isinstance(path, str):
|
||||
raise ValueError(f"Invalid path type: {type(path)}")
|
||||
|
||||
# Check permission
|
||||
if not self.permission_manager.is_path_allowed(path):
|
||||
raise PermissionError(
|
||||
f"Operation '{self.operation}' not allowed for path: {path}"
|
||||
)
|
||||
|
||||
# Call the function
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
@@ -1,984 +0,0 @@
|
||||
"""Tool personality system for organizing development tools based on famous programmers."""
|
||||
|
||||
import os
|
||||
from typing import Set, Dict, List, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolPersonality:
|
||||
"""Represents a programmer personality with tool preferences."""
|
||||
|
||||
name: str
|
||||
programmer: str
|
||||
description: str
|
||||
tools: List[str]
|
||||
environment: Optional[Dict[str, str]] = None
|
||||
philosophy: Optional[str] = None
|
||||
|
||||
def __post_init__(self):
|
||||
"""Validate personality configuration."""
|
||||
if not self.name:
|
||||
raise ValueError("Personality name is required")
|
||||
if not self.tools:
|
||||
raise ValueError("Personality must include at least one tool")
|
||||
|
||||
|
||||
class PersonalityRegistry:
|
||||
"""Registry for tool personalities."""
|
||||
|
||||
_personalities: Dict[str, ToolPersonality] = {}
|
||||
_active_personality: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def register(cls, personality: ToolPersonality) -> None:
|
||||
"""Register a tool personality."""
|
||||
cls._personalities[personality.name] = personality
|
||||
|
||||
@classmethod
|
||||
def get(cls, name: str) -> Optional[ToolPersonality]:
|
||||
"""Get a personality by name."""
|
||||
return cls._personalities.get(name)
|
||||
|
||||
@classmethod
|
||||
def list(cls) -> List[ToolPersonality]:
|
||||
"""List all registered personalities."""
|
||||
return list(cls._personalities.values())
|
||||
|
||||
@classmethod
|
||||
def set_active(cls, name: str) -> None:
|
||||
"""Set the active personality."""
|
||||
if name not in cls._personalities:
|
||||
raise ValueError(f"Personality '{name}' not found")
|
||||
cls._active_personality = name
|
||||
|
||||
@classmethod
|
||||
def get_active(cls) -> Optional[ToolPersonality]:
|
||||
"""Get the active personality."""
|
||||
if cls._active_personality:
|
||||
return cls._personalities.get(cls._active_personality)
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def get_active_tools(cls) -> Set[str]:
|
||||
"""Get the set of tools from the active personality."""
|
||||
personality = cls.get_active()
|
||||
if personality:
|
||||
return set(personality.tools)
|
||||
return set()
|
||||
|
||||
|
||||
# Essential tools that are always available
|
||||
ESSENTIAL_TOOLS = ["read", "write", "edit", "tree", "bash", "think"]
|
||||
|
||||
# Common tool sets for reuse
|
||||
UNIX_TOOLS = ["grep", "find_files", "bash", "process", "diff"]
|
||||
BUILD_TOOLS = ["bash", "npx", "uvx", "process"]
|
||||
VERSION_CONTROL = ["git_search", "diff"]
|
||||
AI_TOOLS = ["agent", "consensus", "critic", "think"]
|
||||
SEARCH_TOOLS = ["search", "symbols", "grep", "git_search"]
|
||||
DATABASE_TOOLS = ["sql_query", "sql_search", "graph_add", "graph_query"]
|
||||
VECTOR_TOOLS = ["vector_index", "vector_search"]
|
||||
|
||||
# 100 Programmer Personalities
|
||||
personalities = [
|
||||
# 1-10: Language Creators
|
||||
ToolPersonality(
|
||||
name="guido",
|
||||
programmer="Guido van Rossum",
|
||||
description="Python's BDFL - readability counts",
|
||||
philosophy="There should be one-- and preferably only one --obvious way to do it.",
|
||||
tools=ESSENTIAL_TOOLS
|
||||
+ ["uvx", "jupyter", "multi_edit", "symbols", "rules"]
|
||||
+ AI_TOOLS
|
||||
+ SEARCH_TOOLS,
|
||||
environment={"PYTHONPATH": ".", "PYTEST_ARGS": "-xvs"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="matz",
|
||||
programmer="Yukihiro Matsumoto",
|
||||
description="Ruby creator - optimize for developer happiness",
|
||||
philosophy="Ruby is designed to make programmers happy.",
|
||||
tools=ESSENTIAL_TOOLS + ["npx", "symbols", "batch", "todo"] + SEARCH_TOOLS,
|
||||
environment={"RUBY_VERSION": "3.0", "BUNDLE_PATH": "vendor/bundle"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="brendan",
|
||||
programmer="Brendan Eich",
|
||||
description="JavaScript creator - dynamic and flexible",
|
||||
philosophy="Always bet on JS.",
|
||||
tools=ESSENTIAL_TOOLS
|
||||
+ ["npx", "watch", "symbols", "todo", "rules"]
|
||||
+ BUILD_TOOLS
|
||||
+ SEARCH_TOOLS,
|
||||
environment={"NODE_ENV": "development", "NPM_CONFIG_LOGLEVEL": "warn"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="dennis",
|
||||
programmer="Dennis Ritchie",
|
||||
description="C creator - close to the metal",
|
||||
philosophy="UNIX is basically a simple operating system, but you have to be a genius to understand the simplicity.",
|
||||
tools=ESSENTIAL_TOOLS + ["symbols", "content_replace"] + UNIX_TOOLS,
|
||||
environment={"CC": "gcc", "CFLAGS": "-Wall -O2"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="bjarne",
|
||||
programmer="Bjarne Stroustrup",
|
||||
description="C++ creator - zero-overhead abstractions",
|
||||
philosophy="C++ is designed to allow you to express ideas.",
|
||||
tools=ESSENTIAL_TOOLS
|
||||
+ ["symbols", "multi_edit", "content_replace"]
|
||||
+ UNIX_TOOLS
|
||||
+ BUILD_TOOLS,
|
||||
environment={"CXX": "g++", "CXXFLAGS": "-std=c++20 -Wall"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="james",
|
||||
programmer="James Gosling",
|
||||
description="Java creator - write once, run anywhere",
|
||||
philosophy="Java is C++ without the guns, knives, and clubs.",
|
||||
tools=ESSENTIAL_TOOLS + ["symbols", "batch", "todo"] + BUILD_TOOLS,
|
||||
environment={"JAVA_HOME": "/usr/lib/jvm/java-11-openjdk"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="anders",
|
||||
programmer="Anders Hejlsberg",
|
||||
description="TypeScript/C# creator - type safety matters",
|
||||
philosophy="TypeScript is JavaScript that scales.",
|
||||
tools=ESSENTIAL_TOOLS
|
||||
+ ["npx", "symbols", "watch", "rules"]
|
||||
+ BUILD_TOOLS
|
||||
+ SEARCH_TOOLS,
|
||||
environment={"TYPESCRIPT_VERSION": "5.0"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="larry",
|
||||
programmer="Larry Wall",
|
||||
description="Perl creator - there's more than one way to do it",
|
||||
philosophy="The three chief virtues of a programmer are laziness, impatience, and hubris.",
|
||||
tools=ESSENTIAL_TOOLS + ["grep", "content_replace", "batch"] + UNIX_TOOLS,
|
||||
environment={"PERL5LIB": "./lib"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="rasmus",
|
||||
programmer="Rasmus Lerdorf",
|
||||
description="PHP creator - pragmatic web development",
|
||||
philosophy="I'm not a real programmer. I throw together things until it works.",
|
||||
tools=ESSENTIAL_TOOLS + ["npx", "sql_query", "watch"] + DATABASE_TOOLS,
|
||||
environment={"PHP_VERSION": "8.0"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="rich",
|
||||
programmer="Rich Hickey",
|
||||
description="Clojure creator - simplicity matters",
|
||||
philosophy="Programming is not about typing... it's about thinking.",
|
||||
tools=ESSENTIAL_TOOLS + ["symbols", "todo", "batch"] + AI_TOOLS,
|
||||
environment={"CLOJURE_VERSION": "1.11"},
|
||||
),
|
||||
# 11-20: Systems & Infrastructure
|
||||
ToolPersonality(
|
||||
name="linus",
|
||||
programmer="Linus Torvalds",
|
||||
description="Linux & Git creator - pragmatic excellence",
|
||||
philosophy="Talk is cheap. Show me the code.",
|
||||
tools=ESSENTIAL_TOOLS
|
||||
+ ["git_search", "diff", "content_replace", "critic"]
|
||||
+ UNIX_TOOLS,
|
||||
environment={"KERNEL_VERSION": "6.0", "GIT_AUTHOR_NAME": "Linus Torvalds"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="rob",
|
||||
programmer="Rob Pike",
|
||||
description="Go creator - simplicity and concurrency",
|
||||
philosophy="A little copying is better than a little dependency.",
|
||||
tools=ESSENTIAL_TOOLS
|
||||
+ ["symbols", "batch", "process"]
|
||||
+ UNIX_TOOLS
|
||||
+ BUILD_TOOLS,
|
||||
environment={"GOPATH": "~/go", "GO111MODULE": "on"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="ken",
|
||||
programmer="Ken Thompson",
|
||||
description="Unix creator - elegant minimalism",
|
||||
philosophy="When in doubt, use brute force.",
|
||||
tools=ESSENTIAL_TOOLS + UNIX_TOOLS,
|
||||
environment={"PATH": "/usr/local/bin:$PATH"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="bill",
|
||||
programmer="Bill Joy",
|
||||
description="Vi creator & BSD contributor",
|
||||
philosophy="The best way to predict the future is to invent it.",
|
||||
tools=ESSENTIAL_TOOLS + ["neovim_edit", "neovim_command"] + UNIX_TOOLS,
|
||||
environment={"EDITOR": "vi"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="richard",
|
||||
programmer="Richard Stallman",
|
||||
description="GNU creator - software freedom",
|
||||
philosophy="Free software is a matter of liberty, not price.",
|
||||
tools=ESSENTIAL_TOOLS + ["content_replace", "batch"] + UNIX_TOOLS,
|
||||
environment={"EDITOR": "emacs"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="brian",
|
||||
programmer="Brian Kernighan",
|
||||
description="AWK co-creator & Unix pioneer",
|
||||
philosophy="Controlling complexity is the essence of computer programming.",
|
||||
tools=ESSENTIAL_TOOLS + ["grep", "content_replace"] + UNIX_TOOLS,
|
||||
environment={"AWK": "gawk"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="donald",
|
||||
programmer="Donald Knuth",
|
||||
description="TeX creator - literate programming",
|
||||
philosophy="Premature optimization is the root of all evil.",
|
||||
tools=ESSENTIAL_TOOLS + ["symbols", "todo", "critic"],
|
||||
environment={"TEXMFHOME": "~/texmf"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="graydon",
|
||||
programmer="Graydon Hoare",
|
||||
description="Rust creator - memory safety without GC",
|
||||
philosophy="Memory safety without garbage collection, concurrency without data races.",
|
||||
tools=ESSENTIAL_TOOLS
|
||||
+ ["symbols", "multi_edit", "critic", "todo"]
|
||||
+ BUILD_TOOLS,
|
||||
environment={"RUST_BACKTRACE": "1", "CARGO_HOME": "~/.cargo"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="ryan",
|
||||
programmer="Ryan Dahl",
|
||||
description="Node.js & Deno creator",
|
||||
philosophy="I/O needs to be done differently.",
|
||||
tools=ESSENTIAL_TOOLS + ["npx", "uvx", "watch", "process"] + BUILD_TOOLS,
|
||||
environment={"DENO_DIR": "~/.deno"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="mitchell",
|
||||
programmer="Mitchell Hashimoto",
|
||||
description="HashiCorp founder - infrastructure as code",
|
||||
philosophy="Automate everything.",
|
||||
tools=ESSENTIAL_TOOLS + ["bash", "process", "watch", "todo"] + BUILD_TOOLS,
|
||||
environment={"TERRAFORM_VERSION": "1.0"},
|
||||
),
|
||||
# 21-30: Web & Frontend
|
||||
ToolPersonality(
|
||||
name="tim",
|
||||
programmer="Tim Berners-Lee",
|
||||
description="WWW inventor - open web",
|
||||
philosophy="The Web is for everyone.",
|
||||
tools=ESSENTIAL_TOOLS + ["npx", "watch", "rules"] + SEARCH_TOOLS,
|
||||
environment={"W3C_VALIDATOR": "true"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="douglas",
|
||||
programmer="Douglas Crockford",
|
||||
description="JSON creator - JavaScript the good parts",
|
||||
philosophy="JavaScript has some extraordinarily good parts.",
|
||||
tools=ESSENTIAL_TOOLS + ["npx", "symbols", "critic"] + SEARCH_TOOLS,
|
||||
environment={"JSLINT": "true"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="john",
|
||||
programmer="John Resig",
|
||||
description="jQuery creator - write less, do more",
|
||||
philosophy="Do more with less code.",
|
||||
tools=ESSENTIAL_TOOLS + ["npx", "watch", "symbols"] + SEARCH_TOOLS,
|
||||
environment={"JQUERY_VERSION": "3.6"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="evan",
|
||||
programmer="Evan You",
|
||||
description="Vue.js creator - progressive framework",
|
||||
philosophy="Approachable, versatile, performant.",
|
||||
tools=ESSENTIAL_TOOLS + ["npx", "watch", "symbols", "todo"] + BUILD_TOOLS,
|
||||
environment={"VUE_VERSION": "3"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="jordan",
|
||||
programmer="Jordan Walke",
|
||||
description="React creator - declarative UIs",
|
||||
philosophy="Learn once, write anywhere.",
|
||||
tools=ESSENTIAL_TOOLS + ["npx", "watch", "symbols", "rules"] + BUILD_TOOLS,
|
||||
environment={"REACT_VERSION": "18"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="jeremy",
|
||||
programmer="Jeremy Ashkenas",
|
||||
description="CoffeeScript & Backbone creator",
|
||||
philosophy="It's just JavaScript.",
|
||||
tools=ESSENTIAL_TOOLS + ["npx", "symbols", "watch"],
|
||||
environment={"COFFEE_VERSION": "2.0"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="david",
|
||||
programmer="David Heinemeier Hansson",
|
||||
description="Rails creator - convention over configuration",
|
||||
philosophy="Optimize for programmer happiness.",
|
||||
tools=ESSENTIAL_TOOLS + ["npx", "sql_query", "watch", "todo"] + DATABASE_TOOLS,
|
||||
environment={"RAILS_ENV": "development"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="taylor",
|
||||
programmer="Taylor Otwell",
|
||||
description="Laravel creator - PHP artisan",
|
||||
philosophy="Love beautiful code? We do too.",
|
||||
tools=ESSENTIAL_TOOLS + ["npx", "sql_query", "watch"] + DATABASE_TOOLS,
|
||||
environment={"LARAVEL_VERSION": "10"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="adrian",
|
||||
programmer="Adrian Holovaty",
|
||||
description="Django co-creator - web framework for perfectionists",
|
||||
philosophy="The web framework for perfectionists with deadlines.",
|
||||
tools=ESSENTIAL_TOOLS + ["uvx", "sql_query", "watch"] + DATABASE_TOOLS,
|
||||
environment={"DJANGO_SETTINGS_MODULE": "settings"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="matt",
|
||||
programmer="Matt Mullenweg",
|
||||
description="WordPress creator - democratize publishing",
|
||||
philosophy="Code is poetry.",
|
||||
tools=ESSENTIAL_TOOLS + ["sql_query", "watch", "rules"] + DATABASE_TOOLS,
|
||||
environment={"WP_DEBUG": "true"},
|
||||
),
|
||||
# 31-40: Database & Data
|
||||
ToolPersonality(
|
||||
name="michael_s",
|
||||
programmer="Michael Stonebraker",
|
||||
description="PostgreSQL creator - ACID matters",
|
||||
philosophy="One size does not fit all in databases.",
|
||||
tools=ESSENTIAL_TOOLS + DATABASE_TOOLS + ["batch", "todo"],
|
||||
environment={"PGDATA": "/var/lib/postgresql/data"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="michael_w",
|
||||
programmer="Michael Widenius",
|
||||
description="MySQL/MariaDB creator",
|
||||
philosophy="A small fast database for the web.",
|
||||
tools=ESSENTIAL_TOOLS + DATABASE_TOOLS + ["watch"],
|
||||
environment={"MYSQL_HOME": "/usr/local/mysql"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="salvatore",
|
||||
programmer="Salvatore Sanfilippo",
|
||||
description="Redis creator - data structures server",
|
||||
philosophy="Simplicity is a great virtue.",
|
||||
tools=ESSENTIAL_TOOLS + ["bash", "watch", "process"] + DATABASE_TOOLS,
|
||||
environment={"REDIS_VERSION": "7.0"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="dwight",
|
||||
programmer="Dwight Merriman",
|
||||
description="MongoDB co-creator - document databases",
|
||||
philosophy="Build the database you want to use.",
|
||||
tools=ESSENTIAL_TOOLS + DATABASE_TOOLS + ["watch", "todo"],
|
||||
environment={"MONGO_VERSION": "6.0"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="edgar",
|
||||
programmer="Edgar F. Codd",
|
||||
description="Relational model inventor",
|
||||
philosophy="Data independence is key.",
|
||||
tools=ESSENTIAL_TOOLS + DATABASE_TOOLS + ["critic"],
|
||||
environment={"SQL_MODE": "ANSI"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="jim_gray",
|
||||
programmer="Jim Gray",
|
||||
description="Transaction processing pioneer",
|
||||
philosophy="The transaction is the unit of work.",
|
||||
tools=ESSENTIAL_TOOLS + DATABASE_TOOLS + ["batch", "critic"],
|
||||
environment={"ISOLATION_LEVEL": "SERIALIZABLE"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="jeff_dean",
|
||||
programmer="Jeff Dean",
|
||||
description="MapReduce & BigTable co-creator",
|
||||
philosophy="Design for planet-scale.",
|
||||
tools=ESSENTIAL_TOOLS + DATABASE_TOOLS + VECTOR_TOOLS + ["batch"],
|
||||
environment={"HADOOP_HOME": "/opt/hadoop"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="sanjay",
|
||||
programmer="Sanjay Ghemawat",
|
||||
description="MapReduce & BigTable co-creator",
|
||||
philosophy="Simple abstractions for complex systems.",
|
||||
tools=ESSENTIAL_TOOLS + DATABASE_TOOLS + ["batch", "process"],
|
||||
environment={"SPARK_HOME": "/opt/spark"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="mike",
|
||||
programmer="Mike Cafarella",
|
||||
description="Hadoop co-creator",
|
||||
philosophy="Storage is cheap, compute is cheap.",
|
||||
tools=ESSENTIAL_TOOLS + DATABASE_TOOLS + ["batch", "process"],
|
||||
environment={"HADOOP_CONF_DIR": "/etc/hadoop"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="matei",
|
||||
programmer="Matei Zaharia",
|
||||
description="Apache Spark creator",
|
||||
philosophy="In-memory computing changes everything.",
|
||||
tools=ESSENTIAL_TOOLS + DATABASE_TOOLS + ["batch", "process", "jupyter"],
|
||||
environment={"SPARK_MASTER": "local[*]"},
|
||||
),
|
||||
# 41-50: AI & Machine Learning
|
||||
ToolPersonality(
|
||||
name="yann",
|
||||
programmer="Yann LeCun",
|
||||
description="Deep learning pioneer - ConvNets",
|
||||
philosophy="AI is not magic; it's just math and data.",
|
||||
tools=ESSENTIAL_TOOLS + AI_TOOLS + VECTOR_TOOLS + ["jupyter", "watch"],
|
||||
environment={"PYTORCH_VERSION": "2.0"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="geoffrey",
|
||||
programmer="Geoffrey Hinton",
|
||||
description="Deep learning godfather",
|
||||
philosophy="The brain has to work with what it's got.",
|
||||
tools=ESSENTIAL_TOOLS + AI_TOOLS + VECTOR_TOOLS + ["jupyter"],
|
||||
environment={"TF_VERSION": "2.13"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="yoshua",
|
||||
programmer="Yoshua Bengio",
|
||||
description="Deep learning pioneer",
|
||||
philosophy="We need to think about AI that helps humanity.",
|
||||
tools=ESSENTIAL_TOOLS + AI_TOOLS + VECTOR_TOOLS + ["jupyter", "batch"],
|
||||
environment={"THEANO_FLAGS": "device=cuda"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="andrew",
|
||||
programmer="Andrew Ng",
|
||||
description="AI educator & Coursera co-founder",
|
||||
philosophy="AI is the new electricity.",
|
||||
tools=ESSENTIAL_TOOLS + AI_TOOLS + ["jupyter", "todo", "watch"],
|
||||
environment={"CUDA_VISIBLE_DEVICES": "0"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="demis",
|
||||
programmer="Demis Hassabis",
|
||||
description="DeepMind co-founder",
|
||||
philosophy="Solve intelligence, use it to solve everything else.",
|
||||
tools=ESSENTIAL_TOOLS + AI_TOOLS + VECTOR_TOOLS + ["agent", "consensus"],
|
||||
environment={"JAX_VERSION": "0.4"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="ilya",
|
||||
programmer="Ilya Sutskever",
|
||||
description="OpenAI co-founder",
|
||||
philosophy="Scale is all you need.",
|
||||
tools=ESSENTIAL_TOOLS + AI_TOOLS + ["agent", "consensus", "critic"],
|
||||
environment={"OPENAI_API_KEY": "sk-..."},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="andrej",
|
||||
programmer="Andrej Karpathy",
|
||||
description="AI educator & Tesla AI director",
|
||||
philosophy="Build it from scratch to understand it.",
|
||||
tools=ESSENTIAL_TOOLS + AI_TOOLS + ["jupyter", "watch", "todo"],
|
||||
environment={"CUDA_HOME": "/usr/local/cuda"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="chris",
|
||||
programmer="Chris Olah",
|
||||
description="AI interpretability researcher",
|
||||
philosophy="Understanding neural networks matters.",
|
||||
tools=ESSENTIAL_TOOLS + AI_TOOLS + VECTOR_TOOLS + ["jupyter", "critic"],
|
||||
environment={"DISTILL_MODE": "interactive"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="francois",
|
||||
programmer="François Chollet",
|
||||
description="Keras creator",
|
||||
philosophy="Deep learning for humans.",
|
||||
tools=ESSENTIAL_TOOLS + AI_TOOLS + ["jupyter", "watch", "todo"],
|
||||
environment={"KERAS_BACKEND": "tensorflow"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="jeremy_howard",
|
||||
programmer="Jeremy Howard",
|
||||
description="fast.ai founder",
|
||||
philosophy="Deep learning should be accessible to all.",
|
||||
tools=ESSENTIAL_TOOLS + AI_TOOLS + ["jupyter", "watch", "rules"],
|
||||
environment={"FASTAI_VERSION": "2.7"},
|
||||
),
|
||||
# 51-60: Security & Cryptography
|
||||
ToolPersonality(
|
||||
name="bruce",
|
||||
programmer="Bruce Schneier",
|
||||
description="Security expert & cryptographer",
|
||||
philosophy="Security is a process, not a product.",
|
||||
tools=ESSENTIAL_TOOLS + ["critic", "symbols", "git_search"] + UNIX_TOOLS,
|
||||
environment={"SECURITY_AUDIT": "true"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="phil",
|
||||
programmer="Phil Zimmermann",
|
||||
description="PGP creator - privacy matters",
|
||||
philosophy="If privacy is outlawed, only outlaws will have privacy.",
|
||||
tools=ESSENTIAL_TOOLS + ["critic", "content_replace"] + UNIX_TOOLS,
|
||||
environment={"GPG_TTY": "$(tty)"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="whitfield",
|
||||
programmer="Whitfield Diffie",
|
||||
description="Public-key cryptography pioneer",
|
||||
philosophy="Privacy is necessary for an open society.",
|
||||
tools=ESSENTIAL_TOOLS + ["critic", "symbols"],
|
||||
environment={"OPENSSL_VERSION": "3.0"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="ralph",
|
||||
programmer="Ralph Merkle",
|
||||
description="Merkle trees inventor",
|
||||
philosophy="Cryptography is about mathematical guarantees.",
|
||||
tools=ESSENTIAL_TOOLS + ["symbols", "critic", "batch"],
|
||||
environment={"HASH_ALGORITHM": "SHA256"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="daniel_b",
|
||||
programmer="Daniel J. Bernstein",
|
||||
description="djb - qmail & Curve25519 creator",
|
||||
philosophy="Security through simplicity.",
|
||||
tools=ESSENTIAL_TOOLS + ["critic", "symbols"] + UNIX_TOOLS,
|
||||
environment={"QMAIL_HOME": "/var/qmail"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="moxie",
|
||||
programmer="Moxie Marlinspike",
|
||||
description="Signal creator - privacy for everyone",
|
||||
philosophy="Making private communication simple.",
|
||||
tools=ESSENTIAL_TOOLS + ["critic", "symbols", "rules"],
|
||||
environment={"SIGNAL_PROTOCOL": "true"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="theo",
|
||||
programmer="Theo de Raadt",
|
||||
description="OpenBSD creator - security by default",
|
||||
philosophy="Shut up and hack.",
|
||||
tools=ESSENTIAL_TOOLS + ["critic", "diff"] + UNIX_TOOLS,
|
||||
environment={"OPENBSD_VERSION": "7.3"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="dan_kaminsky",
|
||||
programmer="Dan Kaminsky",
|
||||
description="DNS security researcher",
|
||||
philosophy="Break it to make it better.",
|
||||
tools=ESSENTIAL_TOOLS + ["critic", "symbols", "process"] + UNIX_TOOLS,
|
||||
environment={"DNSSEC": "true"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="katie",
|
||||
programmer="Katie Moussouris",
|
||||
description="Bug bounty pioneer",
|
||||
philosophy="Hackers are a resource, not a threat.",
|
||||
tools=ESSENTIAL_TOOLS + ["critic", "symbols", "todo"],
|
||||
environment={"BUG_BOUNTY": "enabled"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="matt_blaze",
|
||||
programmer="Matt Blaze",
|
||||
description="Cryptographer & security researcher",
|
||||
philosophy="Crypto is hard to get right.",
|
||||
tools=ESSENTIAL_TOOLS + ["critic", "symbols", "git_search"],
|
||||
environment={"CRYPTO_LIBRARY": "nacl"},
|
||||
),
|
||||
# 61-70: Gaming & Graphics
|
||||
ToolPersonality(
|
||||
name="john_carmack",
|
||||
programmer="John Carmack",
|
||||
description="id Software - Doom & Quake creator",
|
||||
philosophy="Focus is a matter of deciding what things you're not going to do.",
|
||||
tools=ESSENTIAL_TOOLS + ["symbols", "watch", "process"] + BUILD_TOOLS,
|
||||
environment={"OPENGL_VERSION": "4.6"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="sid",
|
||||
programmer="Sid Meier",
|
||||
description="Civilization creator",
|
||||
philosophy="A game is a series of interesting choices.",
|
||||
tools=ESSENTIAL_TOOLS + ["todo", "watch", "process"],
|
||||
environment={"GAME_MODE": "debug"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="shigeru",
|
||||
programmer="Shigeru Miyamoto",
|
||||
description="Mario & Zelda creator",
|
||||
philosophy="A delayed game is eventually good, but a rushed game is forever bad.",
|
||||
tools=ESSENTIAL_TOOLS + ["todo", "watch", "critic"],
|
||||
environment={"NINTENDO_SDK": "true"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="gabe",
|
||||
programmer="Gabe Newell",
|
||||
description="Valve founder - Half-Life & Steam",
|
||||
philosophy="The easiest way to stop piracy is not by putting antipiracy technology to work. It's by giving those people a service that's better than what they're receiving from the pirates.",
|
||||
tools=ESSENTIAL_TOOLS + ["process", "watch", "todo"] + BUILD_TOOLS,
|
||||
environment={"STEAM_RUNTIME": "1"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="markus",
|
||||
programmer="Markus Persson",
|
||||
description="Minecraft creator - Notch",
|
||||
philosophy="Just make games for yourself and try to have fun.",
|
||||
tools=ESSENTIAL_TOOLS + ["watch", "todo", "process"],
|
||||
environment={"LWJGL_VERSION": "3.3"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="jonathan",
|
||||
programmer="Jonathan Blow",
|
||||
description="Braid & The Witness creator",
|
||||
philosophy="Optimize for deep, meaningful experiences.",
|
||||
tools=ESSENTIAL_TOOLS + ["symbols", "critic", "watch"],
|
||||
environment={"JAI_COMPILER": "beta"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="casey",
|
||||
programmer="Casey Muratori",
|
||||
description="Handmade Hero creator",
|
||||
philosophy="Performance matters. Write code from scratch.",
|
||||
tools=ESSENTIAL_TOOLS + ["symbols", "watch", "process", "critic"],
|
||||
environment={"HANDMADE": "true"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="tim_sweeney",
|
||||
programmer="Tim Sweeney",
|
||||
description="Epic Games founder - Unreal Engine",
|
||||
philosophy="The engine is the game.",
|
||||
tools=ESSENTIAL_TOOLS + ["symbols", "watch", "process"] + BUILD_TOOLS,
|
||||
environment={"UNREAL_ENGINE": "5"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="hideo",
|
||||
programmer="Hideo Kojima",
|
||||
description="Metal Gear creator",
|
||||
philosophy="70% of my body is made of movies.",
|
||||
tools=ESSENTIAL_TOOLS + ["todo", "watch", "critic"],
|
||||
environment={"KOJIMA_PRODUCTIONS": "true"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="will",
|
||||
programmer="Will Wright",
|
||||
description="SimCity & The Sims creator",
|
||||
philosophy="Games are a form of communication.",
|
||||
tools=ESSENTIAL_TOOLS + ["todo", "watch", "process"],
|
||||
environment={"SIMULATION_MODE": "debug"},
|
||||
),
|
||||
# 71-80: Open Source Leaders
|
||||
ToolPersonality(
|
||||
name="miguel",
|
||||
programmer="Miguel de Icaza",
|
||||
description="GNOME & Mono creator",
|
||||
philosophy="Open source is about standing on the shoulders of giants.",
|
||||
tools=ESSENTIAL_TOOLS + ["symbols", "todo"] + BUILD_TOOLS,
|
||||
environment={"MONO_VERSION": "6.12"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="nat",
|
||||
programmer="Nat Friedman",
|
||||
description="GitHub CEO & AI entrepreneur",
|
||||
philosophy="Developers are the builders of the digital world.",
|
||||
tools=ESSENTIAL_TOOLS + AI_TOOLS + ["git_search", "todo"],
|
||||
environment={"GITHUB_TOKEN": "ghp_..."},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="patrick",
|
||||
programmer="Patrick Volkerding",
|
||||
description="Slackware creator",
|
||||
philosophy="Keep it simple, keep it stable.",
|
||||
tools=ESSENTIAL_TOOLS + UNIX_TOOLS,
|
||||
environment={"SLACKWARE_VERSION": "15.0"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="ian",
|
||||
programmer="Ian Murdock",
|
||||
description="Debian founder",
|
||||
philosophy="Free software, free society.",
|
||||
tools=ESSENTIAL_TOOLS + UNIX_TOOLS + ["todo"],
|
||||
environment={"DEBIAN_FRONTEND": "noninteractive"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="mark_shuttleworth",
|
||||
programmer="Mark Shuttleworth",
|
||||
description="Ubuntu founder",
|
||||
philosophy="Linux for human beings.",
|
||||
tools=ESSENTIAL_TOOLS + ["todo", "rules"] + BUILD_TOOLS,
|
||||
environment={"UBUNTU_VERSION": "22.04"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="lennart",
|
||||
programmer="Lennart Poettering",
|
||||
description="systemd creator",
|
||||
philosophy="Do one thing and do it well... or do everything.",
|
||||
tools=ESSENTIAL_TOOLS + ["process", "watch"] + UNIX_TOOLS,
|
||||
environment={"SYSTEMD_VERSION": "253"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="bram",
|
||||
programmer="Bram Moolenaar",
|
||||
description="Vim creator",
|
||||
philosophy="The best way to avoid RSI is to not type so much.",
|
||||
tools=ESSENTIAL_TOOLS + ["neovim_edit", "neovim_command", "neovim_session"],
|
||||
environment={"VIM_VERSION": "9.0"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="daniel_r",
|
||||
programmer="Daniel Robbins",
|
||||
description="Gentoo founder",
|
||||
philosophy="Your system, your way.",
|
||||
tools=ESSENTIAL_TOOLS + BUILD_TOOLS + UNIX_TOOLS,
|
||||
environment={"GENTOO_PROFILE": "default/linux/amd64/17.1"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="judd",
|
||||
programmer="Judd Vinet",
|
||||
description="Arch Linux creator",
|
||||
philosophy="Keep it simple.",
|
||||
tools=ESSENTIAL_TOOLS + BUILD_TOOLS + UNIX_TOOLS,
|
||||
environment={"ARCH_VERSION": "rolling"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="fabrice",
|
||||
programmer="Fabrice Bellard",
|
||||
description="QEMU & FFmpeg creator",
|
||||
philosophy="Small, fast, and elegant code.",
|
||||
tools=ESSENTIAL_TOOLS + ["symbols", "process"] + BUILD_TOOLS,
|
||||
environment={"QEMU_VERSION": "8.0"},
|
||||
),
|
||||
# 81-90: Modern Innovators
|
||||
ToolPersonality(
|
||||
name="vitalik",
|
||||
programmer="Vitalik Buterin",
|
||||
description="Ethereum creator",
|
||||
philosophy="Decentralization matters.",
|
||||
tools=ESSENTIAL_TOOLS + ["symbols", "critic", "todo"] + AI_TOOLS,
|
||||
environment={"ETH_NETWORK": "mainnet"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="satoshi",
|
||||
programmer="Satoshi Nakamoto",
|
||||
description="Bitcoin creator",
|
||||
philosophy="Trust in mathematics.",
|
||||
tools=ESSENTIAL_TOOLS + ["critic", "symbols"] + UNIX_TOOLS,
|
||||
environment={"BITCOIN_NETWORK": "mainnet"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="chris_lattner",
|
||||
programmer="Chris Lattner",
|
||||
description="LLVM & Swift creator",
|
||||
philosophy="Compiler infrastructure should be modular.",
|
||||
tools=ESSENTIAL_TOOLS + ["symbols", "multi_edit", "critic"] + BUILD_TOOLS,
|
||||
environment={"LLVM_VERSION": "16"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="joe",
|
||||
programmer="Joe Armstrong",
|
||||
description="Erlang creator",
|
||||
philosophy="Let it crash.",
|
||||
tools=ESSENTIAL_TOOLS + ["process", "watch", "critic"],
|
||||
environment={"ERL_VERSION": "OTP-26"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="jose",
|
||||
programmer="José Valim",
|
||||
description="Elixir creator",
|
||||
philosophy="Productive. Reliable. Fast.",
|
||||
tools=ESSENTIAL_TOOLS + ["watch", "process", "todo"],
|
||||
environment={"ELIXIR_VERSION": "1.15"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="sebastian",
|
||||
programmer="Sebastian Thrun",
|
||||
description="Udacity founder & self-driving car pioneer",
|
||||
philosophy="Education should be accessible to all.",
|
||||
tools=ESSENTIAL_TOOLS + AI_TOOLS + ["jupyter", "watch"],
|
||||
environment={"ROS_VERSION": "noetic"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="palmer",
|
||||
programmer="Palmer Luckey",
|
||||
description="Oculus founder",
|
||||
philosophy="VR is the final medium.",
|
||||
tools=ESSENTIAL_TOOLS + ["watch", "process"] + BUILD_TOOLS,
|
||||
environment={"UNITY_VERSION": "2023.1"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="dylan",
|
||||
programmer="Dylan Field",
|
||||
description="Figma co-founder",
|
||||
philosophy="Design tools should be collaborative.",
|
||||
tools=ESSENTIAL_TOOLS + ["watch", "todo", "rules"],
|
||||
environment={"FIGMA_API": "enabled"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="guillermo",
|
||||
programmer="Guillermo Rauch",
|
||||
description="Vercel founder & Next.js creator",
|
||||
philosophy="Make the Web. Faster.",
|
||||
tools=ESSENTIAL_TOOLS + ["npx", "watch", "rules"] + BUILD_TOOLS,
|
||||
environment={"NEXT_VERSION": "14"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="tom",
|
||||
programmer="Tom Preston-Werner",
|
||||
description="GitHub co-founder & TOML creator",
|
||||
philosophy="Optimize for happiness.",
|
||||
tools=ESSENTIAL_TOOLS + ["git_search", "todo", "rules"],
|
||||
environment={"GITHUB_ACTIONS": "true"},
|
||||
),
|
||||
# 91-100: Special Configurations
|
||||
ToolPersonality(
|
||||
name="fullstack",
|
||||
programmer="Full Stack Developer",
|
||||
description="Every tool for every job",
|
||||
philosophy="Jack of all trades, master of... well, all trades.",
|
||||
tools=list(
|
||||
set(
|
||||
ESSENTIAL_TOOLS
|
||||
+ AI_TOOLS
|
||||
+ SEARCH_TOOLS
|
||||
+ DATABASE_TOOLS
|
||||
+ BUILD_TOOLS
|
||||
+ UNIX_TOOLS
|
||||
+ VECTOR_TOOLS
|
||||
+ [
|
||||
"todo",
|
||||
"rules",
|
||||
"watch",
|
||||
"jupyter",
|
||||
"neovim_edit",
|
||||
"mcp",
|
||||
"consensus",
|
||||
]
|
||||
)
|
||||
),
|
||||
environment={"ALL_TOOLS": "enabled"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="minimal",
|
||||
programmer="Minimalist",
|
||||
description="Just the essentials",
|
||||
philosophy="Less is more.",
|
||||
tools=ESSENTIAL_TOOLS,
|
||||
environment={"MINIMAL_MODE": "true"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="data_scientist",
|
||||
programmer="Data Scientist",
|
||||
description="Analyze all the things",
|
||||
philosophy="In God we trust. All others must bring data.",
|
||||
tools=ESSENTIAL_TOOLS
|
||||
+ ["jupyter", "sql_query", "stats"]
|
||||
+ VECTOR_TOOLS
|
||||
+ AI_TOOLS,
|
||||
environment={"JUPYTER_THEME": "dark"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="devops",
|
||||
programmer="DevOps Engineer",
|
||||
description="Automate everything",
|
||||
philosophy="You build it, you run it.",
|
||||
tools=ESSENTIAL_TOOLS + BUILD_TOOLS + ["process", "watch", "todo"] + UNIX_TOOLS,
|
||||
environment={"CI_CD": "enabled"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="security",
|
||||
programmer="Security Researcher",
|
||||
description="Break it to secure it",
|
||||
philosophy="The only secure system is one that's powered off.",
|
||||
tools=ESSENTIAL_TOOLS + ["critic", "symbols", "git_search"] + UNIX_TOOLS,
|
||||
environment={"SECURITY_MODE": "paranoid"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="academic",
|
||||
programmer="Academic Researcher",
|
||||
description="Publish or perish",
|
||||
philosophy="Standing on the shoulders of giants.",
|
||||
tools=ESSENTIAL_TOOLS + ["jupyter", "todo", "critic"] + AI_TOOLS + SEARCH_TOOLS,
|
||||
environment={"LATEX_ENGINE": "xelatex"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="startup",
|
||||
programmer="Startup Founder",
|
||||
description="Move fast and fix things",
|
||||
philosophy="Done is better than perfect.",
|
||||
tools=ESSENTIAL_TOOLS
|
||||
+ ["todo", "agent", "consensus"]
|
||||
+ BUILD_TOOLS
|
||||
+ DATABASE_TOOLS,
|
||||
environment={"STARTUP_MODE": "hustle"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="enterprise",
|
||||
programmer="Enterprise Developer",
|
||||
description="Process and compliance",
|
||||
philosophy="Nobody ever got fired for buying IBM.",
|
||||
tools=ESSENTIAL_TOOLS + ["todo", "critic", "rules", "stats"] + DATABASE_TOOLS,
|
||||
environment={"COMPLIANCE": "SOC2"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="creative",
|
||||
programmer="Creative Coder",
|
||||
description="Code as art",
|
||||
philosophy="Programming is the art of the possible.",
|
||||
tools=ESSENTIAL_TOOLS + ["watch", "jupyter", "todo"] + AI_TOOLS,
|
||||
environment={"P5_MODE": "global"},
|
||||
),
|
||||
ToolPersonality(
|
||||
name="hanzo",
|
||||
programmer="Hanzo AI Default",
|
||||
description="Balanced productivity and quality",
|
||||
philosophy="The Zen of Model Context Protocol.",
|
||||
tools=ESSENTIAL_TOOLS
|
||||
+ [
|
||||
"agent",
|
||||
"consensus",
|
||||
"critic",
|
||||
"todo",
|
||||
"rules",
|
||||
"symbols",
|
||||
"search",
|
||||
"git_search",
|
||||
"watch",
|
||||
"jupyter",
|
||||
]
|
||||
+ BUILD_TOOLS,
|
||||
environment={"HANZO_MODE": "zen"},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
# Register all personalities
|
||||
def register_default_personalities():
|
||||
"""Register all default tool personalities."""
|
||||
for personality in personalities:
|
||||
PersonalityRegistry.register(personality)
|
||||
|
||||
|
||||
# Ensure agent tool is enabled when API keys are present
|
||||
def ensure_agent_enabled(personality: ToolPersonality) -> ToolPersonality:
|
||||
"""Ensure agent tool is enabled if API keys are present."""
|
||||
api_keys_present = any(
|
||||
os.environ.get(key)
|
||||
for key in [
|
||||
"OPENAI_API_KEY",
|
||||
"ANTHROPIC_API_KEY",
|
||||
"GOOGLE_API_KEY",
|
||||
"HANZO_API_KEY",
|
||||
"GROQ_API_KEY",
|
||||
"TOGETHER_API_KEY",
|
||||
"MISTRAL_API_KEY",
|
||||
"PERPLEXITY_API_KEY",
|
||||
]
|
||||
)
|
||||
|
||||
if api_keys_present and "agent" not in personality.tools:
|
||||
personality.tools.append("agent")
|
||||
if "consensus" not in personality.tools:
|
||||
personality.tools.append("consensus")
|
||||
|
||||
return personality
|
||||
@@ -1,287 +0,0 @@
|
||||
"""Plugin loader for custom user tools."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import inspect
|
||||
import importlib.util
|
||||
from typing import Any, Dict, List, Type, Optional
|
||||
from pathlib import Path
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .base import BaseTool
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolPlugin:
|
||||
"""Represents a loaded tool plugin."""
|
||||
|
||||
name: str
|
||||
tool_class: Type[BaseTool]
|
||||
source_path: Path
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class PluginLoader:
|
||||
"""Loads custom tool plugins from user directories."""
|
||||
|
||||
def __init__(self):
|
||||
self.plugins: Dict[str, ToolPlugin] = {}
|
||||
self.plugin_dirs: List[Path] = []
|
||||
self._setup_plugin_directories()
|
||||
|
||||
def _setup_plugin_directories(self):
|
||||
"""Set up standard plugin directories."""
|
||||
# User's home directory plugins
|
||||
home_plugins = Path.home() / ".hanzo" / "plugins"
|
||||
home_plugins.mkdir(parents=True, exist_ok=True)
|
||||
self.plugin_dirs.append(home_plugins)
|
||||
|
||||
# Project-local plugins
|
||||
project_plugins = Path.cwd() / ".hanzo" / "plugins"
|
||||
if project_plugins.exists():
|
||||
self.plugin_dirs.append(project_plugins)
|
||||
|
||||
# Environment variable for additional paths
|
||||
if custom_paths := os.environ.get("HANZO_PLUGIN_PATH"):
|
||||
for path in custom_paths.split(":"):
|
||||
plugin_dir = Path(path)
|
||||
if plugin_dir.exists():
|
||||
self.plugin_dirs.append(plugin_dir)
|
||||
|
||||
def load_plugins(self) -> Dict[str, ToolPlugin]:
|
||||
"""Load all plugins from configured directories."""
|
||||
for plugin_dir in self.plugin_dirs:
|
||||
if not plugin_dir.exists():
|
||||
continue
|
||||
|
||||
# Look for Python files
|
||||
for py_file in plugin_dir.glob("*.py"):
|
||||
if py_file.name.startswith("_"):
|
||||
continue
|
||||
|
||||
try:
|
||||
self._load_plugin_file(py_file)
|
||||
except Exception as e:
|
||||
print(f"Failed to load plugin {py_file}: {e}")
|
||||
|
||||
# Look for plugin packages
|
||||
for package_dir in plugin_dir.iterdir():
|
||||
if package_dir.is_dir() and (package_dir / "__init__.py").exists():
|
||||
try:
|
||||
self._load_plugin_package(package_dir)
|
||||
except Exception as e:
|
||||
print(f"Failed to load plugin package {package_dir}: {e}")
|
||||
|
||||
return self.plugins
|
||||
|
||||
def _load_plugin_file(self, file_path: Path):
|
||||
"""Load a single plugin file."""
|
||||
# Load the module
|
||||
spec = importlib.util.spec_from_file_location(file_path.stem, file_path)
|
||||
if not spec or not spec.loader:
|
||||
return
|
||||
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[file_path.stem] = module
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
# Find tool classes
|
||||
for _name, obj in inspect.getmembers(module):
|
||||
if (
|
||||
inspect.isclass(obj)
|
||||
and issubclass(obj, BaseTool)
|
||||
and obj != BaseTool
|
||||
and hasattr(obj, "name")
|
||||
):
|
||||
# Load metadata if available
|
||||
metadata = None
|
||||
metadata_file = file_path.with_suffix(".json")
|
||||
if metadata_file.exists():
|
||||
with open(metadata_file) as f:
|
||||
metadata = json.load(f)
|
||||
|
||||
plugin = ToolPlugin(
|
||||
name=obj.name,
|
||||
tool_class=obj,
|
||||
source_path=file_path,
|
||||
metadata=metadata,
|
||||
)
|
||||
self.plugins[obj.name] = plugin
|
||||
|
||||
def _load_plugin_package(self, package_dir: Path):
|
||||
"""Load a plugin package."""
|
||||
# Add parent to path temporarily
|
||||
parent = str(package_dir.parent)
|
||||
if parent not in sys.path:
|
||||
sys.path.insert(0, parent)
|
||||
|
||||
try:
|
||||
# Import the package
|
||||
module = importlib.import_module(package_dir.name)
|
||||
|
||||
# Look for tools
|
||||
if hasattr(module, "TOOLS"):
|
||||
# Package exports TOOLS list
|
||||
for tool_class in module.TOOLS:
|
||||
if issubclass(tool_class, BaseTool):
|
||||
plugin = ToolPlugin(
|
||||
name=tool_class.name,
|
||||
tool_class=tool_class,
|
||||
source_path=package_dir,
|
||||
)
|
||||
self.plugins[tool_class.name] = plugin
|
||||
else:
|
||||
# Search for tool classes
|
||||
for _name, obj in inspect.getmembers(module):
|
||||
if (
|
||||
inspect.isclass(obj)
|
||||
and issubclass(obj, BaseTool)
|
||||
and obj != BaseTool
|
||||
and hasattr(obj, "name")
|
||||
):
|
||||
plugin = ToolPlugin(
|
||||
name=obj.name, tool_class=obj, source_path=package_dir
|
||||
)
|
||||
self.plugins[obj.name] = plugin
|
||||
finally:
|
||||
# Remove from path
|
||||
if parent in sys.path:
|
||||
sys.path.remove(parent)
|
||||
|
||||
def get_tool_class(self, name: str) -> Optional[Type[BaseTool]]:
|
||||
"""Get a tool class by name."""
|
||||
plugin = self.plugins.get(name)
|
||||
return plugin.tool_class if plugin else None
|
||||
|
||||
def list_plugins(self) -> List[str]:
|
||||
"""List all loaded plugin names."""
|
||||
return list(self.plugins.keys())
|
||||
|
||||
|
||||
# Global plugin loader instance
|
||||
_plugin_loader = PluginLoader()
|
||||
|
||||
|
||||
def load_user_plugins() -> Dict[str, ToolPlugin]:
|
||||
"""Load all user plugins."""
|
||||
return _plugin_loader.load_plugins()
|
||||
|
||||
|
||||
def get_plugin_tool(name: str) -> Optional[Type[BaseTool]]:
|
||||
"""Get a plugin tool class by name."""
|
||||
return _plugin_loader.get_tool_class(name)
|
||||
|
||||
|
||||
def list_plugin_tools() -> List[str]:
|
||||
"""List all available plugin tools."""
|
||||
return _plugin_loader.list_plugins()
|
||||
|
||||
|
||||
def create_plugin_template(output_dir: Path, tool_name: str):
|
||||
"""Create a template for a new plugin tool."""
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Create tool file
|
||||
tool_file = output_dir / f"{tool_name}_tool.py"
|
||||
tool_content = f'''"""Custom {tool_name} tool plugin."""
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from typing import Dict, Any
|
||||
|
||||
|
||||
class {tool_name.title()}Tool(BaseTool):
|
||||
"""Custom {tool_name} tool implementation."""
|
||||
|
||||
name = "{tool_name}"
|
||||
description = "Custom {tool_name} tool"
|
||||
|
||||
async def run(self, params: Dict[str, Any], ctx) -> Dict[str, Any]:
|
||||
"""Execute the {tool_name} tool."""
|
||||
# Get parameters
|
||||
action = params.get("action", "default")
|
||||
|
||||
# Implement your tool logic here
|
||||
if action == "default":
|
||||
return {{
|
||||
"status": "success",
|
||||
"message": f"Running {tool_name} tool",
|
||||
"data": {{
|
||||
"params": params
|
||||
}}
|
||||
}}
|
||||
|
||||
# Add more actions as needed
|
||||
elif action == "custom_action":
|
||||
# Your custom logic here
|
||||
pass
|
||||
|
||||
return {{
|
||||
"status": "error",
|
||||
"message": f"Unknown action: {{action}}"
|
||||
}}
|
||||
|
||||
|
||||
# Optional: Export tools explicitly
|
||||
TOOLS = [{tool_name.title()}Tool]
|
||||
'''
|
||||
|
||||
with open(tool_file, "w") as f:
|
||||
f.write(tool_content)
|
||||
|
||||
# Create metadata file
|
||||
metadata_file = output_dir / f"{tool_name}_tool.json"
|
||||
metadata_content = {
|
||||
"name": tool_name,
|
||||
"version": "1.0.0",
|
||||
"author": "Your Name",
|
||||
"description": f"Custom {tool_name} tool",
|
||||
"modes": ["custom"], # Modes this tool should be added to
|
||||
"dependencies": [],
|
||||
"config": {
|
||||
# Tool-specific configuration
|
||||
},
|
||||
}
|
||||
|
||||
with open(metadata_file, "w") as f:
|
||||
json.dump(metadata_content, f, indent=2)
|
||||
|
||||
# Create README
|
||||
readme_file = output_dir / "README.md"
|
||||
readme_content = f"""# {tool_name.title()} Tool Plugin
|
||||
|
||||
Custom tool plugin for Hanzo MCP.
|
||||
|
||||
## Installation
|
||||
|
||||
1. Place this directory in one of:
|
||||
- `~/.hanzo/plugins/`
|
||||
- `./.hanzo/plugins/` (project-specific)
|
||||
- Any path in `HANZO_PLUGIN_PATH` environment variable
|
||||
|
||||
2. The tool will be automatically loaded when Hanzo MCP starts.
|
||||
|
||||
## Usage
|
||||
|
||||
The tool will be available as `{tool_name}` in any mode that includes it.
|
||||
|
||||
## Configuration
|
||||
|
||||
Edit the `{tool_name}_tool.json` file to:
|
||||
- Add the tool to specific modes
|
||||
- Configure tool-specific settings
|
||||
- Specify dependencies
|
||||
|
||||
## Development
|
||||
|
||||
Modify `{tool_name}_tool.py` to implement your custom functionality.
|
||||
"""
|
||||
|
||||
with open(readme_file, "w") as f:
|
||||
f.write(readme_content)
|
||||
|
||||
print(f"Created plugin template in {output_dir}")
|
||||
print(f"Files created:")
|
||||
print(f" - {tool_file}")
|
||||
print(f" - {metadata_file}")
|
||||
print(f" - {readme_file}")
|
||||
@@ -1,279 +0,0 @@
|
||||
"""Comprehensive system and MCP statistics."""
|
||||
|
||||
from typing import Unpack, TypedDict, final, override
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
|
||||
import psutil
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.mcp.mcp_add import McpAddTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.shell.run_background import RunBackgroundTool
|
||||
from hanzo_mcp.tools.database.database_manager import DatabaseManager
|
||||
|
||||
|
||||
class StatsParams(TypedDict, total=False):
|
||||
"""Parameters for stats tool."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@final
|
||||
class StatsTool(BaseTool):
|
||||
"""Tool for showing comprehensive system and MCP statistics."""
|
||||
|
||||
def __init__(self, db_manager: DatabaseManager = None):
|
||||
"""Initialize the stats tool.
|
||||
|
||||
Args:
|
||||
db_manager: Optional database manager for DB stats
|
||||
"""
|
||||
self.db_manager = db_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "stats"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Show comprehensive system and Hanzo AI statistics.
|
||||
|
||||
Displays:
|
||||
- System resources (CPU, memory, disk)
|
||||
- Running processes
|
||||
- Database usage
|
||||
- MCP server status
|
||||
- Tool usage statistics
|
||||
- Warnings for high resource usage
|
||||
|
||||
Example:
|
||||
- stats
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[StatsParams],
|
||||
) -> str:
|
||||
"""Get comprehensive statistics.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Comprehensive statistics
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
output = []
|
||||
warnings = []
|
||||
|
||||
# Header
|
||||
output.append("=== Hanzo AI System Statistics ===")
|
||||
output.append(f"Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||||
output.append("")
|
||||
|
||||
# System Resources
|
||||
output.append("=== System Resources ===")
|
||||
|
||||
# CPU
|
||||
cpu_percent = psutil.cpu_percent(interval=1)
|
||||
cpu_count = psutil.cpu_count()
|
||||
output.append(f"CPU Usage: {cpu_percent}% ({cpu_count} cores)")
|
||||
if cpu_percent > 90:
|
||||
warnings.append(f"⚠️ HIGH CPU USAGE: {cpu_percent}%")
|
||||
|
||||
# Memory
|
||||
memory = psutil.virtual_memory()
|
||||
memory_used_gb = memory.used / (1024**3)
|
||||
memory_total_gb = memory.total / (1024**3)
|
||||
memory_percent = memory.percent
|
||||
output.append(
|
||||
f"Memory: {memory_used_gb:.1f}/{memory_total_gb:.1f} GB ({memory_percent}%)"
|
||||
)
|
||||
if memory_percent > 90:
|
||||
warnings.append(f"⚠️ HIGH MEMORY USAGE: {memory_percent}%")
|
||||
|
||||
# Disk
|
||||
disk = psutil.disk_usage("/")
|
||||
disk_used_gb = disk.used / (1024**3)
|
||||
disk_total_gb = disk.total / (1024**3)
|
||||
disk_percent = disk.percent
|
||||
disk_free_gb = disk.free / (1024**3)
|
||||
output.append(
|
||||
f"Disk: {disk_used_gb:.1f}/{disk_total_gb:.1f} GB ({disk_percent}%)"
|
||||
)
|
||||
output.append(f"Free Space: {disk_free_gb:.1f} GB")
|
||||
if disk_percent > 90:
|
||||
warnings.append(
|
||||
f"⚠️ LOW DISK SPACE: Only {disk_free_gb:.1f} GB free ({100 - disk_percent:.1f}% remaining)"
|
||||
)
|
||||
|
||||
output.append("")
|
||||
|
||||
# Background Processes
|
||||
output.append("=== Background Processes ===")
|
||||
processes = RunBackgroundTool.get_processes()
|
||||
running_count = 0
|
||||
total_memory_mb = 0
|
||||
|
||||
if processes:
|
||||
for proc in processes.values():
|
||||
if proc.is_running():
|
||||
running_count += 1
|
||||
try:
|
||||
ps_proc = psutil.Process(proc.process.pid)
|
||||
memory_mb = ps_proc.memory_info().rss / (1024**2)
|
||||
total_memory_mb += memory_mb
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
output.append(f"Running Processes: {running_count}")
|
||||
output.append(f"Total Memory Usage: {total_memory_mb:.1f} MB")
|
||||
|
||||
# List top processes by memory
|
||||
if running_count > 0:
|
||||
output.append("\nTop Processes:")
|
||||
proc_list = []
|
||||
for proc_id, proc in processes.items():
|
||||
if proc.is_running():
|
||||
try:
|
||||
ps_proc = psutil.Process(proc.process.pid)
|
||||
memory_mb = ps_proc.memory_info().rss / (1024**2)
|
||||
cpu = ps_proc.cpu_percent(interval=0.1)
|
||||
proc_list.append((proc.name, memory_mb, cpu, proc_id))
|
||||
except Exception:
|
||||
proc_list.append((proc.name, 0, 0, proc_id))
|
||||
|
||||
proc_list.sort(key=lambda x: x[1], reverse=True)
|
||||
for name, mem, cpu, pid in proc_list[:5]:
|
||||
output.append(f" - {name} ({pid}): {mem:.1f} MB, {cpu:.1f}% CPU")
|
||||
else:
|
||||
output.append("No background processes running")
|
||||
|
||||
output.append("")
|
||||
|
||||
# Database Usage
|
||||
if self.db_manager:
|
||||
output.append("=== Database Usage ===")
|
||||
db_dir = Path.home() / ".hanzo" / "db"
|
||||
total_db_size = 0
|
||||
|
||||
if db_dir.exists():
|
||||
for db_file in db_dir.rglob("*.db"):
|
||||
size = db_file.stat().st_size
|
||||
total_db_size += size
|
||||
|
||||
output.append(
|
||||
f"Total Database Size: {total_db_size / (1024**2):.1f} MB"
|
||||
)
|
||||
output.append(f"Active Projects: {len(self.db_manager.projects)}")
|
||||
|
||||
# List largest databases
|
||||
db_sizes = []
|
||||
for db_file in db_dir.rglob("*.db"):
|
||||
size = db_file.stat().st_size / (1024**2)
|
||||
if size > 0.1: # Only show DBs > 100KB
|
||||
project = db_file.parent.parent.name
|
||||
db_type = db_file.stem
|
||||
db_sizes.append((project, db_type, size))
|
||||
|
||||
if db_sizes:
|
||||
db_sizes.sort(key=lambda x: x[2], reverse=True)
|
||||
output.append("\nLargest Databases:")
|
||||
for project, db_type, size in db_sizes[:5]:
|
||||
output.append(f" - {project}/{db_type}: {size:.1f} MB")
|
||||
else:
|
||||
output.append("No databases found")
|
||||
|
||||
output.append("")
|
||||
|
||||
# MCP Servers
|
||||
output.append("=== MCP Servers ===")
|
||||
mcp_servers = McpAddTool.get_servers()
|
||||
if mcp_servers:
|
||||
running_mcp = sum(
|
||||
1 for s in mcp_servers.values() if s.get("status") == "running"
|
||||
)
|
||||
total_mcp_tools = sum(len(s.get("tools", [])) for s in mcp_servers.values())
|
||||
|
||||
output.append(f"Total Servers: {len(mcp_servers)}")
|
||||
output.append(f"Running: {running_mcp}")
|
||||
output.append(f"Total Tools Available: {total_mcp_tools}")
|
||||
else:
|
||||
output.append("No MCP servers configured")
|
||||
|
||||
output.append("")
|
||||
|
||||
# Hanzo AI Specifics
|
||||
output.append("=== Hanzo AI ===")
|
||||
|
||||
# Log directory size
|
||||
log_dir = Path.home() / ".hanzo" / "logs"
|
||||
if log_dir.exists():
|
||||
log_size = sum(f.stat().st_size for f in log_dir.rglob("*") if f.is_file())
|
||||
log_count = len(list(log_dir.rglob("*.log")))
|
||||
output.append(f"Log Files: {log_count} ({log_size / (1024**2):.1f} MB)")
|
||||
|
||||
if log_size > 100 * 1024**2: # > 100MB
|
||||
warnings.append(
|
||||
f"⚠️ Large log directory: {log_size / (1024**2):.1f} MB"
|
||||
)
|
||||
|
||||
# Config directory
|
||||
config_dir = Path.home() / ".hanzo" / "mcp"
|
||||
if config_dir.exists():
|
||||
config_count = len(list(config_dir.rglob("*.json")))
|
||||
output.append(f"Config Files: {config_count}")
|
||||
|
||||
# Tool status (if available)
|
||||
# TODO: Track tool usage statistics
|
||||
output.append("\nTool Categories:")
|
||||
output.append(" - File Operations: grep, find_files, read, write, edit")
|
||||
output.append(" - Shell: bash, run_background, processes, pkill")
|
||||
output.append(" - Database: sql_query, graph_query, vector_search")
|
||||
output.append(" - Package Runners: uvx, npx, uvx_background, npx_background")
|
||||
output.append(" - MCP Management: mcp_add, mcp_remove, mcp_stats")
|
||||
|
||||
# Warnings Section
|
||||
if warnings:
|
||||
output.append("\n=== ⚠️ WARNINGS ===")
|
||||
for warning in warnings:
|
||||
output.append(warning)
|
||||
output.append("")
|
||||
|
||||
# Recommendations
|
||||
output.append("=== Recommendations ===")
|
||||
if disk_free_gb < 5:
|
||||
output.append("- Free up disk space (< 5GB remaining)")
|
||||
if memory_percent > 80:
|
||||
output.append("- Close unused applications to free memory")
|
||||
if running_count > 10:
|
||||
output.append("- Consider stopping unused background processes")
|
||||
if log_size > 50 * 1024**2:
|
||||
output.append("- Clean up old log files in ~/.hanzo/logs")
|
||||
|
||||
if not any(
|
||||
[
|
||||
disk_free_gb < 5,
|
||||
memory_percent > 80,
|
||||
running_count > 10,
|
||||
log_size > 50 * 1024**2,
|
||||
]
|
||||
):
|
||||
output.append("✅ System resources are healthy")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,31 +0,0 @@
|
||||
"""Test helper classes for MCP tools testing."""
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
class PaginatedResponseWrapper:
|
||||
"""Wrapper class for paginated responses to support tests."""
|
||||
|
||||
def __init__(self, items=None, next_cursor=None, has_more=False, total_items=None):
|
||||
"""Initialize paginated response."""
|
||||
self.items = items or []
|
||||
self.next_cursor = next_cursor
|
||||
self.has_more = has_more
|
||||
self.total_items = total_items or len(self.items)
|
||||
|
||||
def to_json(self) -> Dict[str, Any]:
|
||||
"""Convert to JSON-serializable dict."""
|
||||
return {
|
||||
"items": self.items,
|
||||
"_meta": {
|
||||
"next_cursor": self.next_cursor,
|
||||
"has_more": self.has_more,
|
||||
"total_items": self.total_items,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# Export a convenience constructor
|
||||
def PaginatedResponse(items=None, next_cursor=None, has_more=False, total_items=None):
|
||||
"""Create a paginated response for testing."""
|
||||
return PaginatedResponseWrapper(items, next_cursor, has_more, total_items)
|
||||
@@ -1,147 +0,0 @@
|
||||
"""Thinking tool implementation.
|
||||
|
||||
This module provides the ThinkingTool for Claude to engage in structured thinking.
|
||||
"""
|
||||
|
||||
from typing import Unpack, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
|
||||
Thought = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="The detailed thought process to record",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class ThinkingToolParams(TypedDict):
|
||||
"""Parameters for the ThinkingTool.
|
||||
|
||||
Attributes:
|
||||
thought: The detailed thought process to record
|
||||
"""
|
||||
|
||||
thought: Thought
|
||||
|
||||
|
||||
@final
|
||||
class ThinkingTool(BaseTool):
|
||||
"""Tool for Claude to engage in structured thinking."""
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name.
|
||||
|
||||
Returns:
|
||||
Tool name
|
||||
"""
|
||||
return "think"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description.
|
||||
|
||||
Returns:
|
||||
Tool description
|
||||
"""
|
||||
return """Use the tool to think about something. It will not obtain new information or make any changes to the repository, but just log the thought. Use it when complex reasoning or brainstorming is needed.
|
||||
Ensure thinking content is concise and accurate, without needing to include code details
|
||||
|
||||
Common use cases:
|
||||
1. When exploring a repository and discovering the source of a bug, call this tool to brainstorm several unique ways of fixing the bug, and assess which change(s) are likely to be simplest and most effective
|
||||
2. After receiving test results, use this tool to brainstorm ways to fix failing tests
|
||||
3. When planning a complex refactoring, use this tool to outline different approaches and their tradeoffs
|
||||
4. When designing a new feature, use this tool to think through architecture decisions and implementation details
|
||||
5. When debugging a complex issue, use this tool to organize your thoughts and hypotheses
|
||||
6. When considering changes to the plan or shifts in thinking that the user has not previously mentioned, consider whether it is necessary to confirm with the user.
|
||||
|
||||
<think_example>
|
||||
Feature Implementation Planning
|
||||
- New code search feature requirements:
|
||||
* Search for code patterns across multiple files
|
||||
* Identify function usages and references
|
||||
* Analyze import relationships
|
||||
* Generate summary of matching patterns
|
||||
- Implementation considerations:
|
||||
* Need to leverage existing search mechanisms
|
||||
* Should use regex for pattern matching
|
||||
* Results need consistent format with other search methods
|
||||
* Must handle large codebases efficiently
|
||||
- Design approach:
|
||||
1. Create new CodeSearcher class that follows existing search patterns
|
||||
2. Implement core pattern matching algorithm
|
||||
3. Add result formatting methods
|
||||
4. Integrate with file traversal system
|
||||
5. Add caching for performance optimization
|
||||
- Testing strategy:
|
||||
* Unit tests for search accuracy
|
||||
* Integration tests with existing components
|
||||
* Performance tests with large codebases
|
||||
</think_example>"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialize the thinking tool."""
|
||||
pass
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[ThinkingToolParams],
|
||||
) -> str:
|
||||
"""Execute the tool with the given parameters.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Tool result
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
thought = params.get("thought")
|
||||
|
||||
# Validate required thought parameter
|
||||
if not thought:
|
||||
await tool_ctx.error(
|
||||
"Parameter 'thought' is required but was None or empty"
|
||||
)
|
||||
return "Error: Parameter 'thought' is required but was None or empty"
|
||||
|
||||
if thought.strip() == "":
|
||||
await tool_ctx.error("Parameter 'thought' cannot be empty")
|
||||
return "Error: Parameter 'thought' cannot be empty"
|
||||
|
||||
# Log the thought but don't take action
|
||||
await tool_ctx.info("Thinking process recorded")
|
||||
|
||||
# Return confirmation
|
||||
return "I've recorded your thinking process. You can continue with your next action based on this analysis."
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register this thinking tool with the MCP server.
|
||||
|
||||
Creates a wrapper function with explicitly defined parameters that match
|
||||
the tool's parameter schema and registers it with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
"""
|
||||
tool_self = self # Create a reference to self for use in the closure
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def think(thought: Thought, ctx: MCPContext) -> str:
|
||||
return await tool_self.call(ctx, thought=thought)
|
||||
@@ -1,149 +0,0 @@
|
||||
"""Disable tools dynamically."""
|
||||
|
||||
from typing import Unpack, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.tool_enable import ToolEnableTool
|
||||
|
||||
ToolName = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Name of the tool to disable (e.g., 'grep', 'vector_search')",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
Persist = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Persist the change to config file",
|
||||
default=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class ToolDisableParams(TypedDict, total=False):
|
||||
"""Parameters for tool disable."""
|
||||
|
||||
tool: str
|
||||
persist: bool
|
||||
|
||||
|
||||
@final
|
||||
class ToolDisableTool(BaseTool):
|
||||
"""Tool for disabling other tools dynamically."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the tool disable tool."""
|
||||
# Ensure states are loaded
|
||||
if not ToolEnableTool._initialized:
|
||||
ToolEnableTool._load_states()
|
||||
ToolEnableTool._initialized = True
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "tool_disable"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Disable tools to prevent their use.
|
||||
|
||||
This allows you to temporarily or permanently disable tools.
|
||||
Useful for testing or when a tool is misbehaving.
|
||||
Changes are persisted by default.
|
||||
|
||||
Critical tools (tool_enable, tool_disable, tool_list) cannot be disabled.
|
||||
|
||||
Examples:
|
||||
- tool_disable --tool vector_search
|
||||
- tool_disable --tool uvx_background
|
||||
- tool_disable --tool grep --no-persist
|
||||
|
||||
Use 'tool_list' to see all available tools and their status.
|
||||
Use 'tool_enable' to re-enable disabled tools.
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[ToolDisableParams],
|
||||
) -> str:
|
||||
"""Disable a tool.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Result of disabling the tool
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
tool_name = params.get("tool")
|
||||
if not tool_name:
|
||||
return "Error: tool name is required"
|
||||
|
||||
persist = params.get("persist", True)
|
||||
|
||||
# Prevent disabling critical tools
|
||||
critical_tools = {"tool_enable", "tool_disable", "tool_list", "stats"}
|
||||
if tool_name in critical_tools:
|
||||
return f"Error: Cannot disable critical tool '{tool_name}'. These tools are required for system management."
|
||||
|
||||
# Check current state
|
||||
was_enabled = ToolEnableTool.is_tool_enabled(tool_name)
|
||||
|
||||
if not was_enabled:
|
||||
return f"Tool '{tool_name}' is already disabled."
|
||||
|
||||
# Disable the tool
|
||||
ToolEnableTool._tool_states[tool_name] = False
|
||||
|
||||
# Persist if requested
|
||||
if persist:
|
||||
ToolEnableTool._save_states()
|
||||
await tool_ctx.info(f"Disabled tool '{tool_name}' (persisted)")
|
||||
else:
|
||||
await tool_ctx.info(f"Disabled tool '{tool_name}' (temporary)")
|
||||
|
||||
output = [
|
||||
f"Successfully disabled tool '{tool_name}'",
|
||||
"",
|
||||
"The tool is now unavailable for use.",
|
||||
f"Use 'tool_enable --tool {tool_name}' to re-enable it.",
|
||||
]
|
||||
|
||||
if not persist:
|
||||
output.append(
|
||||
"\nNote: This change is temporary and will be lost on restart."
|
||||
)
|
||||
|
||||
# Warn about commonly used tools
|
||||
common_tools = {"grep", "read", "write", "bash", "edit"}
|
||||
if tool_name in common_tools:
|
||||
output.append(
|
||||
f"\n⚠️ Warning: '{tool_name}' is a commonly used tool. Disabling it may affect normal operations."
|
||||
)
|
||||
|
||||
# Count disabled tools
|
||||
disabled_count = sum(
|
||||
1 for enabled in ToolEnableTool._tool_states.values() if not enabled
|
||||
)
|
||||
output.append(f"\nTotal disabled tools: {disabled_count}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,181 +0,0 @@
|
||||
"""Enable tools dynamically."""
|
||||
|
||||
import json
|
||||
from typing import Unpack, Annotated, TypedDict, final, override
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
|
||||
ToolName = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Name of the tool to enable (e.g., 'grep', 'vector_search')",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
Persist = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Persist the change to config file",
|
||||
default=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class ToolEnableParams(TypedDict, total=False):
|
||||
"""Parameters for tool enable."""
|
||||
|
||||
tool: str
|
||||
persist: bool
|
||||
|
||||
|
||||
@final
|
||||
class ToolEnableTool(BaseTool):
|
||||
"""Tool for enabling other tools dynamically."""
|
||||
|
||||
# Class variable to track enabled/disabled tools
|
||||
_tool_states = {}
|
||||
_config_file = Path.home() / ".hanzo" / "mcp" / "tool_states.json"
|
||||
_initialized = False
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the tool enable tool."""
|
||||
if not ToolEnableTool._initialized:
|
||||
self._load_states()
|
||||
ToolEnableTool._initialized = True
|
||||
|
||||
@classmethod
|
||||
def _load_states(cls):
|
||||
"""Load tool states from config file."""
|
||||
if cls._config_file.exists():
|
||||
try:
|
||||
with open(cls._config_file, "r") as f:
|
||||
cls._tool_states = json.load(f)
|
||||
except Exception:
|
||||
cls._tool_states = {}
|
||||
else:
|
||||
# Default all tools to enabled
|
||||
cls._tool_states = {}
|
||||
|
||||
@classmethod
|
||||
def _save_states(cls):
|
||||
"""Save tool states to config file."""
|
||||
cls._config_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(cls._config_file, "w") as f:
|
||||
json.dump(cls._tool_states, f, indent=2)
|
||||
|
||||
@classmethod
|
||||
def is_tool_enabled(cls, tool_name: str) -> bool:
|
||||
"""Check if a tool is enabled.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
|
||||
Returns:
|
||||
True if enabled (default), False if explicitly disabled
|
||||
"""
|
||||
# Load states if not initialized
|
||||
if not cls._initialized:
|
||||
cls._load_states()
|
||||
cls._initialized = True
|
||||
|
||||
# Default to enabled if not in states
|
||||
return cls._tool_states.get(tool_name, True)
|
||||
|
||||
@classmethod
|
||||
def get_all_states(cls) -> dict:
|
||||
"""Get all tool states."""
|
||||
if not cls._initialized:
|
||||
cls._load_states()
|
||||
cls._initialized = True
|
||||
return cls._tool_states.copy()
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "tool_enable"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Enable tools that have been disabled.
|
||||
|
||||
This allows you to re-enable tools that were previously disabled.
|
||||
Changes are persisted by default.
|
||||
|
||||
Examples:
|
||||
- tool_enable --tool grep
|
||||
- tool_enable --tool vector_search
|
||||
- tool_enable --tool uvx_background --no-persist
|
||||
|
||||
Use 'tool_list' to see all available tools and their status.
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[ToolEnableParams],
|
||||
) -> str:
|
||||
"""Enable a tool.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Result of enabling the tool
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
tool_name = params.get("tool")
|
||||
if not tool_name:
|
||||
return "Error: tool name is required"
|
||||
|
||||
persist = params.get("persist", True)
|
||||
|
||||
# Check current state
|
||||
was_enabled = self.is_tool_enabled(tool_name)
|
||||
|
||||
if was_enabled:
|
||||
return f"Tool '{tool_name}' is already enabled."
|
||||
|
||||
# Enable the tool
|
||||
self._tool_states[tool_name] = True
|
||||
|
||||
# Persist if requested
|
||||
if persist:
|
||||
self._save_states()
|
||||
await tool_ctx.info(f"Enabled tool '{tool_name}' (persisted)")
|
||||
else:
|
||||
await tool_ctx.info(f"Enabled tool '{tool_name}' (temporary)")
|
||||
|
||||
output = [
|
||||
f"Successfully enabled tool '{tool_name}'",
|
||||
"",
|
||||
"The tool is now available for use.",
|
||||
]
|
||||
|
||||
if not persist:
|
||||
output.append("Note: This change is temporary and will be lost on restart.")
|
||||
|
||||
# Count enabled/disabled tools
|
||||
disabled_count = sum(1 for enabled in self._tool_states.values() if not enabled)
|
||||
if disabled_count > 0:
|
||||
output.append(f"\nCurrently disabled tools: {disabled_count}")
|
||||
output.append("Use 'tool_list --disabled' to see them.")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,271 +0,0 @@
|
||||
"""List all available tools and their status."""
|
||||
|
||||
from typing import Unpack, Optional, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.tool_enable import ToolEnableTool
|
||||
|
||||
ShowDisabled = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Show only disabled tools",
|
||||
default=False,
|
||||
),
|
||||
]
|
||||
|
||||
ShowEnabled = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Show only enabled tools",
|
||||
default=False,
|
||||
),
|
||||
]
|
||||
|
||||
Category = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Filter by category (filesystem, shell, database, etc.)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class ToolListParams(TypedDict, total=False):
|
||||
"""Parameters for tool list."""
|
||||
|
||||
show_disabled: bool
|
||||
show_enabled: bool
|
||||
category: Optional[str]
|
||||
|
||||
|
||||
@final
|
||||
class ToolListTool(BaseTool):
|
||||
"""Tool for listing all available tools and their status."""
|
||||
|
||||
# Tool information organized by category
|
||||
TOOL_INFO = {
|
||||
"filesystem": [
|
||||
("read", "Read contents of files"),
|
||||
("write", "Write contents to files"),
|
||||
("edit", "Edit specific parts of files"),
|
||||
("multi_edit", "Make multiple edits to a file"),
|
||||
("tree", "Directory tree visualization (Unix-style)"),
|
||||
("find", "Find text in files (rg/ag/ack/grep)"),
|
||||
("symbols", "Code symbols search with tree-sitter"),
|
||||
("search", "Search (parallel grep/symbols/vector/git)"),
|
||||
("git_search", "Search git history"),
|
||||
("glob", "Find files by name pattern"),
|
||||
("content_replace", "Replace content across files"),
|
||||
],
|
||||
"shell": [
|
||||
("run_command", "Execute shell commands (--background option)"),
|
||||
("streaming_command", "Run commands with disk-based output streaming"),
|
||||
("processes", "List background processes"),
|
||||
("pkill", "Kill background processes"),
|
||||
("logs", "View process logs"),
|
||||
("uvx", "Run Python packages (--background option)"),
|
||||
("npx", "Run Node.js packages (--background option)"),
|
||||
],
|
||||
"database": [
|
||||
("sql", "SQLite operations (query/search/schema/stats)"),
|
||||
("graph", "Graph database (query/add/remove/search/stats)"),
|
||||
("vector", "Semantic search (search/index/stats/clear)"),
|
||||
],
|
||||
"ai": [
|
||||
("llm", "LLM interface (query/consensus/list/models/enable/disable)"),
|
||||
("agent", "AI agents (run/start/call/stop/list with A2A support)"),
|
||||
("swarm", "Parallel agent execution across multiple files"),
|
||||
(
|
||||
"hierarchical_swarm",
|
||||
"Hierarchical agent teams with Claude Code integration",
|
||||
),
|
||||
("mcp", "MCP servers (list/add/remove/enable/disable/restart)"),
|
||||
],
|
||||
"config": [
|
||||
("config", "Git-style configuration (get/set/list/toggle)"),
|
||||
("tool_enable", "Enable tools"),
|
||||
("tool_disable", "Disable tools"),
|
||||
("tool_list", "List all tools (this tool)"),
|
||||
],
|
||||
"productivity": [
|
||||
("todo", "Todo management (list/add/update/remove/clear)"),
|
||||
("jupyter", "Jupyter notebooks (read/edit/create/delete/execute)"),
|
||||
("think", "Structured thinking space"),
|
||||
],
|
||||
"system": [
|
||||
("stats", "System and resource statistics"),
|
||||
("batch", "Run multiple tools in parallel"),
|
||||
],
|
||||
"legacy": [
|
||||
("directory_tree", "Legacy: Use 'tree' instead"),
|
||||
("grep", "Legacy: Use 'find' instead"),
|
||||
("grep_ast", "Legacy: Use 'symbols' instead"),
|
||||
("batch_search", "Legacy: Use 'search' instead"),
|
||||
("find_files", "Legacy: Use 'glob' instead"),
|
||||
("run_background", "Legacy: Use 'run_command --background'"),
|
||||
("uvx_background", "Legacy: Use 'uvx --background'"),
|
||||
("npx_background", "Legacy: Use 'npx --background'"),
|
||||
("sql_query", "Legacy: Use 'sql' instead"),
|
||||
("sql_search", "Legacy: Use 'sql --action search'"),
|
||||
("sql_stats", "Legacy: Use 'sql --action stats'"),
|
||||
("graph_add", "Legacy: Use 'graph --action add'"),
|
||||
("graph_remove", "Legacy: Use 'graph --action remove'"),
|
||||
("graph_query", "Legacy: Use 'graph' instead"),
|
||||
("graph_search", "Legacy: Use 'graph --action search'"),
|
||||
("graph_stats", "Legacy: Use 'graph --action stats'"),
|
||||
("vector_index", "Legacy: Use 'vector --action index'"),
|
||||
("vector_search", "Legacy: Use 'vector' instead"),
|
||||
("dispatch_agent", "Legacy: Use 'agent' instead"),
|
||||
("todo_read", "Legacy: Use 'todo' instead"),
|
||||
("todo_write", "Legacy: Use 'todo --action add/update'"),
|
||||
("notebook_read", "Legacy: Use 'jupyter' instead"),
|
||||
("notebook_edit", "Legacy: Use 'jupyter --action edit'"),
|
||||
],
|
||||
}
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the tool list tool."""
|
||||
pass
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "tool_list"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """List all available tools and their current status.
|
||||
|
||||
Shows:
|
||||
- Tool names and descriptions
|
||||
- Whether each tool is enabled or disabled
|
||||
- Tools organized by category
|
||||
|
||||
Examples:
|
||||
- tool_list # Show all tools
|
||||
- tool_list --show-disabled # Show only disabled tools
|
||||
- tool_list --show-enabled # Show only enabled tools
|
||||
- tool_list --category shell # Show only shell tools
|
||||
|
||||
Use 'tool_enable' and 'tool_disable' to change tool status.
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[ToolListParams],
|
||||
) -> str:
|
||||
"""List all tools.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
List of tools and their status
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
show_disabled = params.get("show_disabled", False)
|
||||
show_enabled = params.get("show_enabled", False)
|
||||
category_filter = params.get("category")
|
||||
|
||||
# Get all tool states
|
||||
all_states = ToolEnableTool.get_all_states()
|
||||
|
||||
output = []
|
||||
|
||||
# Header
|
||||
if show_disabled:
|
||||
output.append("=== Disabled Tools ===")
|
||||
elif show_enabled:
|
||||
output.append("=== Enabled Tools ===")
|
||||
else:
|
||||
output.append("=== All Available Tools ===")
|
||||
|
||||
if category_filter:
|
||||
output.append(f"Category: {category_filter}")
|
||||
|
||||
output.append("")
|
||||
|
||||
# Count statistics
|
||||
total_tools = 0
|
||||
disabled_count = 0
|
||||
shown_count = 0
|
||||
|
||||
# Iterate through categories
|
||||
categories = (
|
||||
[category_filter]
|
||||
if category_filter and category_filter in self.TOOL_INFO
|
||||
else self.TOOL_INFO.keys()
|
||||
)
|
||||
|
||||
for category in categories:
|
||||
if category not in self.TOOL_INFO:
|
||||
continue
|
||||
|
||||
category_tools = self.TOOL_INFO[category]
|
||||
category_shown = []
|
||||
|
||||
for tool_name, description in category_tools:
|
||||
total_tools += 1
|
||||
is_enabled = ToolEnableTool.is_tool_enabled(tool_name)
|
||||
|
||||
if not is_enabled:
|
||||
disabled_count += 1
|
||||
|
||||
# Apply filters
|
||||
if show_disabled and is_enabled:
|
||||
continue
|
||||
if show_enabled and not is_enabled:
|
||||
continue
|
||||
|
||||
status = "✅" if is_enabled else "❌"
|
||||
category_shown.append((tool_name, description, status))
|
||||
shown_count += 1
|
||||
|
||||
# Show category if it has tools
|
||||
if category_shown:
|
||||
output.append(f"=== {category.title()} Tools ===")
|
||||
|
||||
# Find max tool name length for alignment
|
||||
max_name_len = max(len(name) for name, _, _ in category_shown)
|
||||
|
||||
for tool_name, description, status in category_shown:
|
||||
output.append(
|
||||
f"{status} {tool_name.ljust(max_name_len)} - {description}"
|
||||
)
|
||||
|
||||
output.append("")
|
||||
|
||||
# Summary
|
||||
if not show_disabled and not show_enabled:
|
||||
output.append("=== Summary ===")
|
||||
output.append(f"Total tools: {total_tools}")
|
||||
output.append(f"Enabled: {total_tools - disabled_count}")
|
||||
output.append(f"Disabled: {disabled_count}")
|
||||
else:
|
||||
output.append(f"Showing {shown_count} tool(s)")
|
||||
|
||||
if disabled_count > 0 and not show_disabled:
|
||||
output.append("\nUse 'tool_list --show-disabled' to see disabled tools.")
|
||||
output.append("Use 'tool_enable --tool <name>' to enable a tool.")
|
||||
|
||||
if show_disabled:
|
||||
output.append("\nUse 'tool_enable --tool <name>' to enable these tools.")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,101 +0,0 @@
|
||||
"""Response truncation utilities for MCP tools.
|
||||
|
||||
This module provides utilities to ensure MCP tool responses don't exceed token limits.
|
||||
"""
|
||||
|
||||
import tiktoken
|
||||
|
||||
|
||||
def estimate_tokens(text: str, model: str = "gpt-4") -> int:
|
||||
"""Estimate the number of tokens in a text string.
|
||||
|
||||
Args:
|
||||
text: The text to estimate tokens for
|
||||
model: The model to use for token estimation (default: gpt-4)
|
||||
|
||||
Returns:
|
||||
Estimated number of tokens
|
||||
"""
|
||||
try:
|
||||
# Try to get the encoding for the specific model
|
||||
encoding = tiktoken.encoding_for_model(model)
|
||||
except KeyError:
|
||||
# Fall back to cl100k_base which is used by newer models
|
||||
encoding = tiktoken.get_encoding("cl100k_base")
|
||||
|
||||
return len(encoding.encode(text))
|
||||
|
||||
|
||||
def truncate_response(
|
||||
response: str,
|
||||
max_tokens: int = 20000,
|
||||
truncation_message: str = "\n\n[Response truncated due to length. Please use pagination, filtering, or limit parameters to see more.]",
|
||||
) -> str:
|
||||
"""Truncate a response to fit within token limits.
|
||||
|
||||
Args:
|
||||
response: The response text to truncate
|
||||
max_tokens: Maximum number of tokens allowed (default: 20000)
|
||||
truncation_message: Message to append when truncating
|
||||
|
||||
Returns:
|
||||
Truncated response if needed, original response otherwise
|
||||
"""
|
||||
# Quick check - if response is short, no need to count tokens
|
||||
if len(response) < max_tokens * 2: # Rough estimate: 1 token ≈ 2-4 chars
|
||||
return response
|
||||
|
||||
# Estimate tokens
|
||||
token_count = estimate_tokens(response)
|
||||
|
||||
# If within limit, return as-is
|
||||
if token_count <= max_tokens:
|
||||
return response
|
||||
|
||||
# Need to truncate
|
||||
# Binary search to find the right truncation point
|
||||
left, right = 0, len(response)
|
||||
truncation_msg_tokens = estimate_tokens(truncation_message)
|
||||
target_tokens = max_tokens - truncation_msg_tokens
|
||||
|
||||
while left < right - 1:
|
||||
mid = (left + right) // 2
|
||||
mid_tokens = estimate_tokens(response[:mid])
|
||||
|
||||
if mid_tokens <= target_tokens:
|
||||
left = mid
|
||||
else:
|
||||
right = mid
|
||||
|
||||
# Find a good break point (newline or space)
|
||||
truncate_at = left
|
||||
for i in range(min(100, left), -1, -1):
|
||||
if response[left - i] in "\n ":
|
||||
truncate_at = left - i
|
||||
break
|
||||
|
||||
return response[:truncate_at] + truncation_message
|
||||
|
||||
|
||||
def truncate_lines(
|
||||
response: str,
|
||||
max_lines: int = 1000,
|
||||
truncation_message: str = "\n\n[Response truncated to {max_lines} lines. Please use pagination or filtering to see more.]",
|
||||
) -> str:
|
||||
"""Truncate a response by number of lines.
|
||||
|
||||
Args:
|
||||
response: The response text to truncate
|
||||
max_lines: Maximum number of lines allowed (default: 1000)
|
||||
truncation_message: Message template to append when truncating
|
||||
|
||||
Returns:
|
||||
Truncated response if needed, original response otherwise
|
||||
"""
|
||||
lines = response.split("\n")
|
||||
|
||||
if len(lines) <= max_lines:
|
||||
return response
|
||||
|
||||
truncated = "\n".join(lines[:max_lines])
|
||||
return truncated + truncation_message.format(max_lines=max_lines)
|
||||
@@ -1,62 +0,0 @@
|
||||
"""Parameter validation utilities for Hanzo AI tools.
|
||||
|
||||
This module provides utilities for validating parameters in tool functions.
|
||||
"""
|
||||
|
||||
from typing import TypeVar, final
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@final
|
||||
class ValidationResult:
|
||||
"""Result of a parameter validation."""
|
||||
|
||||
def __init__(self, is_valid: bool, error_message: str = ""):
|
||||
"""Initialize a validation result.
|
||||
|
||||
Args:
|
||||
is_valid: Whether the parameter is valid
|
||||
error_message: Optional error message for invalid parameters
|
||||
"""
|
||||
self.is_valid: bool = is_valid
|
||||
self.error_message: str = error_message
|
||||
|
||||
@property
|
||||
def is_error(self) -> bool:
|
||||
"""Check if the validation resulted in an error.
|
||||
|
||||
Returns:
|
||||
True if there was a validation error, False otherwise
|
||||
"""
|
||||
return not self.is_valid
|
||||
|
||||
|
||||
def validate_path_parameter(
|
||||
path: str | None, parameter_name: str = "path"
|
||||
) -> ValidationResult:
|
||||
"""Validate a path parameter.
|
||||
|
||||
Args:
|
||||
path: The path parameter to validate
|
||||
parameter_name: The name of the parameter (for error messages)
|
||||
|
||||
Returns:
|
||||
A ValidationResult indicating whether the parameter is valid
|
||||
"""
|
||||
# Check for None
|
||||
if path is None:
|
||||
return ValidationResult(
|
||||
is_valid=False,
|
||||
error_message=f"Path parameter '{parameter_name}' is required but was None",
|
||||
)
|
||||
|
||||
# Check for empty path
|
||||
if path.strip() == "":
|
||||
return ValidationResult(
|
||||
is_valid=False,
|
||||
error_message=f"Path parameter '{parameter_name}' is required but was empty string",
|
||||
)
|
||||
|
||||
# Path is valid
|
||||
return ValidationResult(is_valid=True)
|
||||
@@ -1,12 +0,0 @@
|
||||
"""Configuration tools for Hanzo AI."""
|
||||
|
||||
from hanzo_mcp.tools.config.mode_tool import mode_tool
|
||||
from hanzo_mcp.tools.config.config_tool import ConfigTool
|
||||
from hanzo_mcp.tools.config.index_config import IndexScope, IndexConfig
|
||||
|
||||
__all__ = [
|
||||
"ConfigTool",
|
||||
"IndexConfig",
|
||||
"IndexScope",
|
||||
"mode_tool",
|
||||
]
|
||||
@@ -1,337 +0,0 @@
|
||||
"""Configuration tool for Hanzo AI.
|
||||
|
||||
Git-style config tool for managing settings.
|
||||
"""
|
||||
|
||||
from typing import Unpack, Optional, Annotated, TypedDict, final, override
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.config import load_settings, save_settings
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.config.index_config import IndexScope, IndexConfig
|
||||
|
||||
# Parameter types
|
||||
Action = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Action: get (default), set, list, toggle",
|
||||
default="get",
|
||||
),
|
||||
]
|
||||
|
||||
Key = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Configuration key (e.g., tools.write.enabled, enabled_tools.write, index.scope)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Value = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Configuration value",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Scope = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Config scope: local (project) or global",
|
||||
default="local",
|
||||
),
|
||||
]
|
||||
|
||||
ConfigPath = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Path for project-specific config",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class ConfigParams(TypedDict, total=False):
|
||||
"""Parameters for config tool."""
|
||||
|
||||
action: str
|
||||
key: Optional[str]
|
||||
value: Optional[str]
|
||||
scope: str
|
||||
path: Optional[str]
|
||||
|
||||
|
||||
def _parse_bool(value: str) -> Optional[bool]:
|
||||
v = value.strip().lower()
|
||||
if v in {"true", "1", "yes", "on"}:
|
||||
return True
|
||||
if v in {"false", "0", "no", "off"}:
|
||||
return False
|
||||
return None
|
||||
|
||||
|
||||
@final
|
||||
class ConfigTool(BaseTool):
|
||||
"""Git-style configuration management tool."""
|
||||
|
||||
def __init__(self, permission_manager: PermissionManager):
|
||||
"""Initialize config tool."""
|
||||
super().__init__(permission_manager)
|
||||
self.index_config = IndexConfig()
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "config"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Git-style configuration. Actions: get (default), set, list, toggle.
|
||||
|
||||
Usage:
|
||||
config index.scope
|
||||
config --action set index.scope project
|
||||
config --action set tools.write.enabled false
|
||||
config --action list
|
||||
config --action toggle index.scope --path ./project"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[ConfigParams],
|
||||
) -> str:
|
||||
"""Execute config operation."""
|
||||
tool_ctx = self.create_tool_context(ctx)
|
||||
|
||||
# Extract parameters
|
||||
action = params.get("action", "get")
|
||||
key = params.get("key")
|
||||
value = params.get("value")
|
||||
scope = params.get("scope", "local")
|
||||
path = params.get("path")
|
||||
|
||||
# Route to handler
|
||||
if action == "get":
|
||||
return await self._handle_get(key, scope, path, tool_ctx)
|
||||
elif action == "set":
|
||||
return await self._handle_set(key, value, scope, path, tool_ctx)
|
||||
elif action == "list":
|
||||
return await self._handle_list(scope, path, tool_ctx)
|
||||
elif action == "toggle":
|
||||
return await self._handle_toggle(key, scope, path, tool_ctx)
|
||||
else:
|
||||
return f"Error: Unknown action '{action}'. Valid actions: get, set, list, toggle"
|
||||
|
||||
async def _handle_get(
|
||||
self, key: Optional[str], scope: str, path: Optional[str], tool_ctx
|
||||
) -> str:
|
||||
"""Get configuration value."""
|
||||
if not key:
|
||||
return "Error: key required for get action"
|
||||
|
||||
# Handle index scope
|
||||
if key == "index.scope":
|
||||
current_scope = self.index_config.get_scope(path)
|
||||
return f"index.scope={current_scope.value}"
|
||||
|
||||
# tools.<name>.enabled → enabled_tools lookup
|
||||
if key.startswith("tools.") and key.endswith(".enabled"):
|
||||
parts = key.split(".")
|
||||
if len(parts) == 3:
|
||||
tool_name = parts[1]
|
||||
settings = load_settings(project_dir=path if scope == "local" else None)
|
||||
return f"{key}={settings.is_tool_enabled(tool_name)}"
|
||||
|
||||
# enabled_tools.<name>
|
||||
if key.startswith("enabled_tools."):
|
||||
tool_name = key.split(".", 1)[1]
|
||||
settings = load_settings(project_dir=path if scope == "local" else None)
|
||||
val = settings.enabled_tools.get(tool_name)
|
||||
return f"{key}={val if val is not None else 'unset'}"
|
||||
|
||||
# Indexing (legacy) per-tool setting: <tool>.enabled
|
||||
if "." in key:
|
||||
tool, setting = key.split(".", 1)
|
||||
if setting == "enabled":
|
||||
enabled = self.index_config.is_indexing_enabled(tool)
|
||||
return f"{key}={enabled}"
|
||||
|
||||
return f"Unknown key: {key}"
|
||||
|
||||
def _save_project_settings(self, settings, project_dir: Optional[str]) -> Path:
|
||||
"""Save to project config if path provided; else global."""
|
||||
if project_dir:
|
||||
project_path = Path(project_dir)
|
||||
project_path.mkdir(parents=True, exist_ok=True)
|
||||
cfg = project_path / ".hanzo-mcp.json"
|
||||
cfg.write_text(__import__("json").dumps(settings.__dict__ if hasattr(settings, "__dict__") else {}, indent=2))
|
||||
return cfg
|
||||
# Fallback to global handler
|
||||
return save_settings(settings, global_config=True)
|
||||
|
||||
async def _handle_set(
|
||||
self,
|
||||
key: Optional[str],
|
||||
value: Optional[str],
|
||||
scope: str,
|
||||
path: Optional[str],
|
||||
tool_ctx,
|
||||
) -> str:
|
||||
"""Set configuration value."""
|
||||
if not key:
|
||||
return "Error: key required for set action"
|
||||
if value is None:
|
||||
return "Error: value required for set action"
|
||||
|
||||
# Handle index scope
|
||||
if key == "index.scope":
|
||||
try:
|
||||
new_scope = IndexScope(value)
|
||||
self.index_config.set_scope(
|
||||
new_scope, path if scope == "local" else None
|
||||
)
|
||||
return f"Set {key}={value} ({'project' if path else 'global'})"
|
||||
except ValueError:
|
||||
return f"Error: Invalid scope value '{value}'. Valid: project, global, auto"
|
||||
|
||||
# tools.<name>.enabled → enabled_tools mapping
|
||||
if key.startswith("tools.") and key.endswith(".enabled"):
|
||||
parts = key.split(".")
|
||||
if len(parts) == 3:
|
||||
tool_name = parts[1]
|
||||
parsed = _parse_bool(value)
|
||||
if parsed is None:
|
||||
return "Error: value must be boolean (true/false)"
|
||||
settings = load_settings(project_dir=path if scope == "local" else None)
|
||||
et = dict(settings.enabled_tools)
|
||||
et[tool_name] = parsed
|
||||
settings.enabled_tools = et
|
||||
# Save
|
||||
if scope == "local" and path:
|
||||
# Write a project .hanzo-mcp.json adjacent to the path
|
||||
# Note: save_settings(local) saves to CWD; we target specific path here
|
||||
out = self._save_project_settings(settings, path)
|
||||
return f"Set {key}={parsed} (project: {out})"
|
||||
else:
|
||||
out = save_settings(settings, global_config=True)
|
||||
return f"Set {key}={parsed} (global: {out})"
|
||||
|
||||
# enabled_tools.<name>
|
||||
if key.startswith("enabled_tools."):
|
||||
tool_name = key.split(".", 1)[1]
|
||||
parsed = _parse_bool(value)
|
||||
if parsed is None:
|
||||
return "Error: value must be boolean (true/false)"
|
||||
settings = load_settings(project_dir=path if scope == "local" else None)
|
||||
et = dict(settings.enabled_tools)
|
||||
et[tool_name] = parsed
|
||||
settings.enabled_tools = et
|
||||
if scope == "local" and path:
|
||||
out = self._save_project_settings(settings, path)
|
||||
return f"Set {key}={parsed} (project: {out})"
|
||||
else:
|
||||
out = save_settings(settings, global_config=True)
|
||||
return f"Set {key}={parsed} (global: {out})"
|
||||
|
||||
# Indexing (legacy) per-tool setting: <tool>.enabled (search indexers)
|
||||
if "." in key:
|
||||
tool, setting = key.split(".", 1)
|
||||
if setting == "enabled":
|
||||
parsed = _parse_bool(value)
|
||||
if parsed is None:
|
||||
return "Error: value must be boolean (true/false)"
|
||||
self.index_config.set_indexing_enabled(tool, parsed)
|
||||
return f"Set {key}={parsed}"
|
||||
|
||||
return f"Unknown key: {key}"
|
||||
|
||||
async def _handle_list(self, scope: str, path: Optional[str], tool_ctx) -> str:
|
||||
"""List all configuration."""
|
||||
status = self.index_config.get_status()
|
||||
|
||||
output = ["=== Configuration ==="]
|
||||
output.append(f"\nDefault scope: {status['default_scope']}")
|
||||
|
||||
if path:
|
||||
current_scope = self.index_config.get_scope(path)
|
||||
output.append(f"Current path scope: {current_scope.value}")
|
||||
|
||||
output.append(f"\nProjects with custom config: {status['project_count']}")
|
||||
|
||||
output.append("\nTool settings (indexing):")
|
||||
for tool, settings in status["tools"].items():
|
||||
output.append(f" {tool}:")
|
||||
output.append(f" enabled: {settings['enabled']}")
|
||||
output.append(f" per_project: {settings['per_project']}")
|
||||
|
||||
# Also show enabled_tools snapshot
|
||||
settings_snapshot = load_settings(project_dir=path if scope == "local" else None)
|
||||
output.append("\nEnabled tools (execution):")
|
||||
for tool_name, enabled in sorted(settings_snapshot.enabled_tools.items()):
|
||||
output.append(f" {tool_name}: {enabled}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
async def _handle_toggle(
|
||||
self, key: Optional[str], scope: str, path: Optional[str], tool_ctx
|
||||
) -> str:
|
||||
"""Toggle configuration value."""
|
||||
if not key:
|
||||
return "Error: key required for toggle action"
|
||||
|
||||
# Handle index scope toggle
|
||||
if key == "index.scope":
|
||||
new_scope = self.index_config.toggle_scope(
|
||||
path if scope == "local" else None
|
||||
)
|
||||
return f"Toggled index.scope to {new_scope.value}"
|
||||
|
||||
# Handle execution tool enable/disable: tools.<name>.enabled or enabled_tools.<name>
|
||||
if key.startswith("tools.") and key.endswith(".enabled"):
|
||||
parts = key.split(".")
|
||||
if len(parts) == 3:
|
||||
tool_name = parts[1]
|
||||
settings = load_settings(project_dir=path if scope == "local" else None)
|
||||
current = bool(settings.enabled_tools.get(tool_name, True))
|
||||
settings.enabled_tools[tool_name] = not current
|
||||
if scope == "local" and path:
|
||||
out = self._save_project_settings(settings, path)
|
||||
return f"Toggled {key} to {not current} (project: {out})"
|
||||
else:
|
||||
out = save_settings(settings, global_config=True)
|
||||
return f"Toggled {key} to {not current} (global: {out})"
|
||||
|
||||
if key.startswith("enabled_tools."):
|
||||
tool_name = key.split(".", 1)[1]
|
||||
settings = load_settings(project_dir=path if scope == "local" else None)
|
||||
current = bool(settings.enabled_tools.get(tool_name, True))
|
||||
settings.enabled_tools[tool_name] = not current
|
||||
if scope == "local" and path:
|
||||
out = self._save_project_settings(settings, path)
|
||||
return f"Toggled {key} to {not current} (project: {out})"
|
||||
else:
|
||||
out = save_settings(settings, global_config=True)
|
||||
return f"Toggled {key} to {not current} (global: {out})"
|
||||
|
||||
# Handle indexing toggles (legacy)
|
||||
if "." in key:
|
||||
tool, setting = key.split(".", 1)
|
||||
if setting == "enabled":
|
||||
current = self.index_config.is_indexing_enabled(tool)
|
||||
self.index_config.set_indexing_enabled(tool, not current)
|
||||
return f"Toggled {key} to {not current}"
|
||||
|
||||
return f"Cannot toggle key: {key}"
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,179 +0,0 @@
|
||||
"""Index configuration for per-project vs global indexing.
|
||||
|
||||
This module manages indexing configuration for different scopes.
|
||||
"""
|
||||
|
||||
import json
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, Optional
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class IndexScope(Enum):
|
||||
"""Indexing scope options."""
|
||||
|
||||
PROJECT = "project" # Per-project indexing
|
||||
GLOBAL = "global" # Global indexing
|
||||
AUTO = "auto" # Auto-detect based on git root
|
||||
|
||||
|
||||
class IndexConfig:
|
||||
"""Manages indexing configuration."""
|
||||
|
||||
def __init__(self, config_dir: Optional[Path] = None):
|
||||
"""Initialize index configuration."""
|
||||
self.config_dir = config_dir or Path.home() / ".hanzo" / "mcp"
|
||||
self.config_file = self.config_dir / "index_config.json"
|
||||
self._config = self._load_config()
|
||||
|
||||
def _load_config(self) -> Dict[str, Any]:
|
||||
"""Load configuration from disk."""
|
||||
if self.config_file.exists():
|
||||
try:
|
||||
with open(self.config_file, "r") as f:
|
||||
return json.load(f)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Default configuration
|
||||
return {
|
||||
"default_scope": IndexScope.AUTO.value,
|
||||
"project_configs": {},
|
||||
"global_index_paths": [],
|
||||
"index_settings": {
|
||||
"vector": {
|
||||
"enabled": True,
|
||||
"auto_index": True,
|
||||
"include_git_history": True,
|
||||
},
|
||||
"symbols": {
|
||||
"enabled": True,
|
||||
"auto_index": False,
|
||||
},
|
||||
"sql": {
|
||||
"enabled": True,
|
||||
"per_project": True,
|
||||
},
|
||||
"graph": {
|
||||
"enabled": True,
|
||||
"per_project": True,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
def save_config(self) -> None:
|
||||
"""Save configuration to disk."""
|
||||
self.config_dir.mkdir(parents=True, exist_ok=True)
|
||||
with open(self.config_file, "w") as f:
|
||||
json.dump(self._config, f, indent=2)
|
||||
|
||||
def get_scope(self, path: Optional[str] = None) -> IndexScope:
|
||||
"""Get indexing scope for a path."""
|
||||
if not path:
|
||||
return IndexScope(self._config["default_scope"])
|
||||
|
||||
# Check project-specific config
|
||||
project_root = self._find_project_root(path)
|
||||
if project_root:
|
||||
project_config = self._config["project_configs"].get(str(project_root))
|
||||
if project_config and "scope" in project_config:
|
||||
return IndexScope(project_config["scope"])
|
||||
|
||||
# Use default
|
||||
scope = IndexScope(self._config["default_scope"])
|
||||
|
||||
# Handle auto mode
|
||||
if scope == IndexScope.AUTO:
|
||||
if project_root:
|
||||
return IndexScope.PROJECT
|
||||
else:
|
||||
return IndexScope.GLOBAL
|
||||
|
||||
return scope
|
||||
|
||||
def set_scope(self, scope: IndexScope, path: Optional[str] = None) -> None:
|
||||
"""Set indexing scope."""
|
||||
if path:
|
||||
# Set for specific project
|
||||
project_root = self._find_project_root(path)
|
||||
if project_root:
|
||||
if str(project_root) not in self._config["project_configs"]:
|
||||
self._config["project_configs"][str(project_root)] = {}
|
||||
self._config["project_configs"][str(project_root)][
|
||||
"scope"
|
||||
] = scope.value
|
||||
else:
|
||||
# Set global default
|
||||
self._config["default_scope"] = scope.value
|
||||
|
||||
self.save_config()
|
||||
|
||||
def get_index_path(self, tool: str, path: Optional[str] = None) -> Path:
|
||||
"""Get index path for a tool and location."""
|
||||
scope = self.get_scope(path)
|
||||
|
||||
if scope == IndexScope.PROJECT and path:
|
||||
project_root = self._find_project_root(path)
|
||||
if project_root:
|
||||
return Path(project_root) / ".hanzo" / "index" / tool
|
||||
|
||||
# Global index
|
||||
return self.config_dir / "index" / tool
|
||||
|
||||
def is_indexing_enabled(self, tool: str) -> bool:
|
||||
"""Check if indexing is enabled for a tool."""
|
||||
return self._config["index_settings"].get(tool, {}).get("enabled", True)
|
||||
|
||||
def set_indexing_enabled(self, tool: str, enabled: bool) -> None:
|
||||
"""Enable/disable indexing for a tool."""
|
||||
if tool not in self._config["index_settings"]:
|
||||
self._config["index_settings"][tool] = {}
|
||||
self._config["index_settings"][tool]["enabled"] = enabled
|
||||
self.save_config()
|
||||
|
||||
def toggle_scope(self, path: Optional[str] = None) -> IndexScope:
|
||||
"""Toggle between project and global scope."""
|
||||
current = self.get_scope(path)
|
||||
|
||||
if current == IndexScope.PROJECT:
|
||||
new_scope = IndexScope.GLOBAL
|
||||
elif current == IndexScope.GLOBAL:
|
||||
new_scope = IndexScope.PROJECT
|
||||
else: # AUTO
|
||||
# Determine what auto resolves to and toggle
|
||||
if path and self._find_project_root(path):
|
||||
new_scope = IndexScope.GLOBAL
|
||||
else:
|
||||
new_scope = IndexScope.PROJECT
|
||||
|
||||
self.set_scope(new_scope, path)
|
||||
return new_scope
|
||||
|
||||
def _find_project_root(self, path: str) -> Optional[Path]:
|
||||
"""Find project root (git root or similar)."""
|
||||
current = Path(path).resolve()
|
||||
|
||||
# Walk up looking for markers
|
||||
markers = [".git", ".hg", "pyproject.toml", "package.json", "Cargo.toml"]
|
||||
|
||||
while current != current.parent:
|
||||
for marker in markers:
|
||||
if (current / marker).exists():
|
||||
return current
|
||||
current = current.parent
|
||||
|
||||
return None
|
||||
|
||||
def get_status(self) -> Dict[str, Any]:
|
||||
"""Get current configuration status."""
|
||||
return {
|
||||
"default_scope": self._config["default_scope"],
|
||||
"project_count": len(self._config["project_configs"]),
|
||||
"tools": {
|
||||
tool: {
|
||||
"enabled": settings.get("enabled", True),
|
||||
"per_project": settings.get("per_project", True),
|
||||
}
|
||||
for tool, settings in self._config["index_settings"].items()
|
||||
},
|
||||
}
|
||||
@@ -1,329 +0,0 @@
|
||||
"""Tool for managing development modes with programmer personalities."""
|
||||
|
||||
from typing import Optional, override
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.mode import ModeRegistry, register_default_modes
|
||||
|
||||
|
||||
class ModeTool(BaseTool):
|
||||
"""Tool for managing development modes."""
|
||||
|
||||
name = "mode"
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the mode tool."""
|
||||
super().__init__()
|
||||
# Register default modes on initialization
|
||||
register_default_modes()
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Manage development modes (programmer personalities). Actions: list (default), activate, show, current.
|
||||
|
||||
Usage:
|
||||
mode
|
||||
mode --action list
|
||||
mode --action activate guido
|
||||
mode --action show linus
|
||||
mode --action current"""
|
||||
|
||||
@override
|
||||
async def run(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
action: str = "list",
|
||||
name: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Manage development modes.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
action: Action to perform (list, activate, show, current)
|
||||
name: Mode name (for activate/show actions)
|
||||
|
||||
Returns:
|
||||
Action result
|
||||
"""
|
||||
if action == "list":
|
||||
modes = ModeRegistry.list()
|
||||
if not modes:
|
||||
return "No modes registered"
|
||||
|
||||
output = ["Available development modes (100 programmer personalities):"]
|
||||
active = ModeRegistry.get_active()
|
||||
|
||||
# Group modes by category
|
||||
categories = {
|
||||
"Language Creators": [
|
||||
"guido",
|
||||
"matz",
|
||||
"brendan",
|
||||
"dennis",
|
||||
"bjarne",
|
||||
"james",
|
||||
"anders",
|
||||
"larry",
|
||||
"rasmus",
|
||||
"rich",
|
||||
],
|
||||
"Systems & Infrastructure": [
|
||||
"linus",
|
||||
"rob",
|
||||
"ken",
|
||||
"bill",
|
||||
"richard",
|
||||
"brian",
|
||||
"donald",
|
||||
"graydon",
|
||||
"ryan",
|
||||
"mitchell",
|
||||
],
|
||||
"Web & Frontend": [
|
||||
"tim",
|
||||
"douglas",
|
||||
"john",
|
||||
"evan",
|
||||
"jordan",
|
||||
"jeremy",
|
||||
"david",
|
||||
"taylor",
|
||||
"adrian",
|
||||
"matt",
|
||||
],
|
||||
"Database & Data": [
|
||||
"michael_s",
|
||||
"michael_w",
|
||||
"salvatore",
|
||||
"dwight",
|
||||
"edgar",
|
||||
"jim_gray",
|
||||
"jeff_dean",
|
||||
"sanjay",
|
||||
"mike",
|
||||
"matei",
|
||||
],
|
||||
"AI & Machine Learning": [
|
||||
"yann",
|
||||
"geoffrey",
|
||||
"yoshua",
|
||||
"andrew",
|
||||
"demis",
|
||||
"ilya",
|
||||
"andrej",
|
||||
"chris",
|
||||
"francois",
|
||||
"jeremy_howard",
|
||||
],
|
||||
"Security & Cryptography": [
|
||||
"bruce",
|
||||
"phil",
|
||||
"whitfield",
|
||||
"ralph",
|
||||
"daniel_b",
|
||||
"moxie",
|
||||
"theo",
|
||||
"dan_kaminsky",
|
||||
"katie",
|
||||
"matt_blaze",
|
||||
],
|
||||
"Gaming & Graphics": [
|
||||
"john_carmack",
|
||||
"sid",
|
||||
"shigeru",
|
||||
"gabe",
|
||||
"markus",
|
||||
"jonathan",
|
||||
"casey",
|
||||
"tim_sweeney",
|
||||
"hideo",
|
||||
"will",
|
||||
],
|
||||
"Open Source Leaders": [
|
||||
"miguel",
|
||||
"nat",
|
||||
"patrick",
|
||||
"ian",
|
||||
"mark_shuttleworth",
|
||||
"lennart",
|
||||
"bram",
|
||||
"daniel_r",
|
||||
"judd",
|
||||
"fabrice",
|
||||
],
|
||||
"Modern Innovators": [
|
||||
"vitalik",
|
||||
"satoshi",
|
||||
"chris_lattner",
|
||||
"joe",
|
||||
"jose",
|
||||
"sebastian",
|
||||
"palmer",
|
||||
"dylan",
|
||||
"guillermo",
|
||||
"tom",
|
||||
],
|
||||
"Special Configurations": [
|
||||
"fullstack",
|
||||
"minimal",
|
||||
"data_scientist",
|
||||
"devops",
|
||||
"security",
|
||||
"academic",
|
||||
"startup",
|
||||
"enterprise",
|
||||
"creative",
|
||||
"hanzo",
|
||||
],
|
||||
}
|
||||
|
||||
for category, mode_names in categories.items():
|
||||
output.append(f"\n{category}:")
|
||||
for mode_name in mode_names:
|
||||
mode = next((m for m in modes if m.name == mode_name), None)
|
||||
if mode:
|
||||
marker = (
|
||||
" (active)" if active and active.name == mode.name else ""
|
||||
)
|
||||
output.append(
|
||||
f" {mode.name}{marker}: {mode.programmer} - {mode.description}"
|
||||
)
|
||||
|
||||
output.append("\nUse 'mode --action activate <name>' to activate a mode")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
elif action == "activate":
|
||||
if not name:
|
||||
return "Error: Mode name required for activate action"
|
||||
|
||||
try:
|
||||
ModeRegistry.set_active(name)
|
||||
mode = ModeRegistry.get(name)
|
||||
|
||||
output = [f"Activated mode: {mode.name}"]
|
||||
output.append(f"Programmer: {mode.programmer}")
|
||||
output.append(f"Description: {mode.description}")
|
||||
if mode.philosophy:
|
||||
output.append(f"Philosophy: {mode.philosophy}")
|
||||
output.append(f"\nEnabled tools ({len(mode.tools)}):")
|
||||
|
||||
# Group tools by category
|
||||
core_tools = []
|
||||
package_tools = []
|
||||
ai_tools = []
|
||||
search_tools = []
|
||||
other_tools = []
|
||||
|
||||
for tool in sorted(mode.tools):
|
||||
if tool in [
|
||||
"read",
|
||||
"write",
|
||||
"edit",
|
||||
"multi_edit",
|
||||
"bash",
|
||||
"tree",
|
||||
"grep",
|
||||
]:
|
||||
core_tools.append(tool)
|
||||
elif tool in ["npx", "uvx", "pip", "cargo", "gem"]:
|
||||
package_tools.append(tool)
|
||||
elif tool in ["agent", "consensus", "critic", "think"]:
|
||||
ai_tools.append(tool)
|
||||
elif tool in ["search", "symbols", "git_search"]:
|
||||
search_tools.append(tool)
|
||||
else:
|
||||
other_tools.append(tool)
|
||||
|
||||
if core_tools:
|
||||
output.append(f" Core: {', '.join(core_tools)}")
|
||||
if package_tools:
|
||||
output.append(f" Package managers: {', '.join(package_tools)}")
|
||||
if ai_tools:
|
||||
output.append(f" AI tools: {', '.join(ai_tools)}")
|
||||
if search_tools:
|
||||
output.append(f" Search: {', '.join(search_tools)}")
|
||||
if other_tools:
|
||||
output.append(f" Specialized: {', '.join(other_tools)}")
|
||||
|
||||
if mode.environment:
|
||||
output.append("\nEnvironment variables:")
|
||||
for key, value in mode.environment.items():
|
||||
output.append(f" {key}={value}")
|
||||
|
||||
output.append(
|
||||
"\nNote: Restart MCP session for changes to take full effect"
|
||||
)
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except ValueError as e:
|
||||
return str(e)
|
||||
|
||||
elif action == "show":
|
||||
if not name:
|
||||
return "Error: Mode name required for show action"
|
||||
|
||||
mode = ModeRegistry.get(name)
|
||||
if not mode:
|
||||
return f"Mode '{name}' not found"
|
||||
|
||||
output = [f"Mode: {mode.name}"]
|
||||
output.append(f"Programmer: {mode.programmer}")
|
||||
output.append(f"Description: {mode.description}")
|
||||
if mode.philosophy:
|
||||
output.append(f"Philosophy: {mode.philosophy}")
|
||||
output.append(f"\nTools ({len(mode.tools)}):")
|
||||
|
||||
for tool in sorted(mode.tools):
|
||||
output.append(f" - {tool}")
|
||||
|
||||
if mode.environment:
|
||||
output.append("\nEnvironment:")
|
||||
for key, value in mode.environment.items():
|
||||
output.append(f" {key}={value}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
elif action == "current":
|
||||
active = ModeRegistry.get_active()
|
||||
if not active:
|
||||
return "No mode currently active\nUse 'mode --action activate <name>' to activate one"
|
||||
|
||||
output = [f"Current mode: {active.name}"]
|
||||
output.append(f"Programmer: {active.programmer}")
|
||||
output.append(f"Description: {active.description}")
|
||||
if active.philosophy:
|
||||
output.append(f"Philosophy: {active.philosophy}")
|
||||
output.append(f"Enabled tools: {len(active.tools)}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
else:
|
||||
return f"Unknown action: {action}. Use 'list', 'activate', 'show', or 'current'"
|
||||
|
||||
def register(self, server: FastMCP) -> None:
|
||||
"""Register the tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@server.tool(name=self.name, description=self.description)
|
||||
async def mode_handler(
|
||||
ctx: MCPContext, action: str = "list", name: Optional[str] = None
|
||||
) -> str:
|
||||
"""Handle mode tool calls."""
|
||||
return await tool_self.run(ctx, action=action, name=name)
|
||||
|
||||
async def call(self, ctx: MCPContext, **params) -> str:
|
||||
"""Call the tool with arguments."""
|
||||
return await self.run(
|
||||
ctx, action=params.get("action", "list"), name=params.get("name")
|
||||
)
|
||||
|
||||
|
||||
# Create tool instance
|
||||
mode_tool = ModeTool()
|
||||
@@ -1,74 +0,0 @@
|
||||
"""Database tools for Hanzo AI.
|
||||
|
||||
This package provides tools for working with embedded SQLite databases
|
||||
and graph databases in projects.
|
||||
"""
|
||||
|
||||
from mcp.server import FastMCP
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
from .graph_add import GraphAddTool
|
||||
|
||||
# Import database tools
|
||||
from .sql_query import SqlQueryTool
|
||||
from .sql_stats import SqlStatsTool
|
||||
from .sql_search import SqlSearchTool
|
||||
from .graph_query import GraphQueryTool
|
||||
from .graph_stats import GraphStatsTool
|
||||
from .graph_remove import GraphRemoveTool
|
||||
from .graph_search import GraphSearchTool
|
||||
from .database_manager import DatabaseManager
|
||||
|
||||
__all__ = [
|
||||
"register_database_tools",
|
||||
"DatabaseManager",
|
||||
"SqlQueryTool",
|
||||
"SqlSearchTool",
|
||||
"SqlStatsTool",
|
||||
"GraphAddTool",
|
||||
"GraphRemoveTool",
|
||||
"GraphQueryTool",
|
||||
"GraphSearchTool",
|
||||
"GraphStatsTool",
|
||||
]
|
||||
|
||||
|
||||
def register_database_tools(
|
||||
mcp_server: FastMCP,
|
||||
permission_manager: PermissionManager,
|
||||
db_manager: DatabaseManager | None = None,
|
||||
) -> list[BaseTool]:
|
||||
"""Register database tools with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
permission_manager: Permission manager for access control
|
||||
db_manager: Optional database manager instance
|
||||
|
||||
Returns:
|
||||
List of registered tools
|
||||
"""
|
||||
# Create database manager if not provided
|
||||
if db_manager is None:
|
||||
db_manager = DatabaseManager(permission_manager)
|
||||
|
||||
# Create tool instances
|
||||
tools = [
|
||||
SqlQueryTool(permission_manager, db_manager),
|
||||
SqlSearchTool(permission_manager, db_manager),
|
||||
SqlStatsTool(permission_manager, db_manager),
|
||||
GraphAddTool(permission_manager, db_manager),
|
||||
GraphRemoveTool(permission_manager, db_manager),
|
||||
GraphQueryTool(permission_manager, db_manager),
|
||||
GraphSearchTool(permission_manager, db_manager),
|
||||
GraphStatsTool(permission_manager, db_manager),
|
||||
]
|
||||
|
||||
# Register with MCP server
|
||||
from hanzo_mcp.tools.common.base import ToolRegistry
|
||||
|
||||
ToolRegistry.register_tools(mcp_server, tools)
|
||||
|
||||
return tools
|
||||
@@ -1,248 +0,0 @@
|
||||
"""Database manager for project-specific SQLite and graph databases."""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
from typing import Dict, List, Optional
|
||||
from pathlib import Path
|
||||
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
|
||||
class ProjectDatabase:
|
||||
"""Manages SQLite and graph databases for a project."""
|
||||
|
||||
def __init__(self, project_path: str):
|
||||
self.project_path = Path(project_path)
|
||||
self.db_dir = self.project_path / ".hanzo" / "db"
|
||||
self.db_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# SQLite database path
|
||||
self.sqlite_path = self.db_dir / "project.db"
|
||||
self.graph_path = self.db_dir / "graph.db"
|
||||
|
||||
# Initialize databases
|
||||
self._init_sqlite()
|
||||
self._init_graph()
|
||||
|
||||
# Keep graph in memory for performance
|
||||
self.graph_conn = sqlite3.connect(":memory:")
|
||||
self._init_graph_schema(self.graph_conn)
|
||||
self._load_graph_from_disk()
|
||||
|
||||
def _init_sqlite(self):
|
||||
"""Initialize SQLite database with common tables."""
|
||||
conn = sqlite3.connect(self.sqlite_path)
|
||||
try:
|
||||
# Create metadata table
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS metadata (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Create files table
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS files (
|
||||
path TEXT PRIMARY KEY,
|
||||
content TEXT,
|
||||
size INTEGER,
|
||||
modified_at TIMESTAMP,
|
||||
hash TEXT,
|
||||
metadata TEXT
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Create symbols table
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS symbols (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
file_path TEXT,
|
||||
name TEXT,
|
||||
type TEXT,
|
||||
line_start INTEGER,
|
||||
line_end INTEGER,
|
||||
scope TEXT,
|
||||
signature TEXT,
|
||||
FOREIGN KEY (file_path) REFERENCES files(path)
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Create index for fast searches
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_files_path ON files(path)")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_symbols_name ON symbols(name)")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_symbols_type ON symbols(type)")
|
||||
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _init_graph(self):
|
||||
"""Initialize graph database on disk."""
|
||||
conn = sqlite3.connect(self.graph_path)
|
||||
try:
|
||||
self._init_graph_schema(conn)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _init_graph_schema(self, conn: sqlite3.Connection):
|
||||
"""Initialize graph database schema."""
|
||||
# Nodes table
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS nodes (
|
||||
id TEXT PRIMARY KEY,
|
||||
type TEXT NOT NULL,
|
||||
properties TEXT,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Edges table
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS edges (
|
||||
source TEXT NOT NULL,
|
||||
target TEXT NOT NULL,
|
||||
relationship TEXT NOT NULL,
|
||||
weight REAL DEFAULT 1.0,
|
||||
properties TEXT,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
PRIMARY KEY (source, target, relationship),
|
||||
FOREIGN KEY (source) REFERENCES nodes(id),
|
||||
FOREIGN KEY (target) REFERENCES nodes(id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Indexes for graph traversal
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_edges_source ON edges(source)")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_edges_target ON edges(target)")
|
||||
conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_edges_relationship ON edges(relationship)"
|
||||
)
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_nodes_type ON nodes(type)")
|
||||
|
||||
def _load_graph_from_disk(self):
|
||||
"""Load graph from disk into memory."""
|
||||
disk_conn = sqlite3.connect(self.graph_path)
|
||||
try:
|
||||
# Copy nodes
|
||||
nodes = disk_conn.execute("SELECT * FROM nodes").fetchall()
|
||||
self.graph_conn.executemany(
|
||||
"INSERT OR REPLACE INTO nodes VALUES (?, ?, ?, ?)", nodes
|
||||
)
|
||||
|
||||
# Copy edges
|
||||
edges = disk_conn.execute("SELECT * FROM edges").fetchall()
|
||||
self.graph_conn.executemany(
|
||||
"INSERT OR REPLACE INTO edges VALUES (?, ?, ?, ?, ?, ?)", edges
|
||||
)
|
||||
|
||||
self.graph_conn.commit()
|
||||
finally:
|
||||
disk_conn.close()
|
||||
|
||||
def _save_graph_to_disk(self):
|
||||
"""Save in-memory graph to disk."""
|
||||
disk_conn = sqlite3.connect(self.graph_path)
|
||||
try:
|
||||
# Clear existing data
|
||||
disk_conn.execute("DELETE FROM edges")
|
||||
disk_conn.execute("DELETE FROM nodes")
|
||||
|
||||
# Copy nodes
|
||||
nodes = self.graph_conn.execute("SELECT * FROM nodes").fetchall()
|
||||
disk_conn.executemany("INSERT INTO nodes VALUES (?, ?, ?, ?)", nodes)
|
||||
|
||||
# Copy edges
|
||||
edges = self.graph_conn.execute("SELECT * FROM edges").fetchall()
|
||||
disk_conn.executemany("INSERT INTO edges VALUES (?, ?, ?, ?, ?, ?)", edges)
|
||||
|
||||
disk_conn.commit()
|
||||
finally:
|
||||
disk_conn.close()
|
||||
|
||||
def get_sqlite_connection(self) -> sqlite3.Connection:
|
||||
"""Get SQLite connection."""
|
||||
conn = sqlite3.connect(self.sqlite_path)
|
||||
conn.row_factory = sqlite3.Row
|
||||
return conn
|
||||
|
||||
def get_graph_connection(self) -> sqlite3.Connection:
|
||||
"""Get in-memory graph connection."""
|
||||
return self.graph_conn
|
||||
|
||||
def close(self):
|
||||
"""Close connections and save graph to disk."""
|
||||
self._save_graph_to_disk()
|
||||
self.graph_conn.close()
|
||||
|
||||
|
||||
class DatabaseManager:
|
||||
"""Manages databases for multiple projects."""
|
||||
|
||||
def __init__(self, permission_manager: PermissionManager):
|
||||
self.permission_manager = permission_manager
|
||||
self.projects: Dict[str, ProjectDatabase] = {}
|
||||
self.search_paths: List[str] = []
|
||||
|
||||
def add_search_path(self, path: str):
|
||||
"""Add a path to search for projects."""
|
||||
if path not in self.search_paths:
|
||||
self.search_paths.append(path)
|
||||
|
||||
def get_project_db(self, project_path: str) -> ProjectDatabase:
|
||||
"""Get or create database for a project."""
|
||||
project_path = os.path.abspath(project_path)
|
||||
|
||||
# Check permissions
|
||||
if not self.permission_manager.has_permission(project_path):
|
||||
raise PermissionError(f"No permission to access: {project_path}")
|
||||
|
||||
# Create database if not exists
|
||||
if project_path not in self.projects:
|
||||
self.projects[project_path] = ProjectDatabase(project_path)
|
||||
|
||||
return self.projects[project_path]
|
||||
|
||||
def get_project_for_path(self, file_path: str) -> Optional[ProjectDatabase]:
|
||||
"""Find the project database for a given file path."""
|
||||
file_path = os.path.abspath(file_path)
|
||||
|
||||
# Check if file is in a known project
|
||||
for project_path in self.projects:
|
||||
if file_path.startswith(project_path):
|
||||
return self.projects[project_path]
|
||||
|
||||
# Search up the directory tree for a project
|
||||
current = Path(file_path)
|
||||
if current.is_file():
|
||||
current = current.parent
|
||||
|
||||
while current != current.parent:
|
||||
# Check for project markers
|
||||
if (current / ".git").exists() or (current / "LLM.md").exists():
|
||||
return self.get_project_db(str(current))
|
||||
current = current.parent
|
||||
|
||||
# No project found, use the directory of the file
|
||||
if Path(file_path).is_file():
|
||||
return self.get_project_db(str(Path(file_path).parent))
|
||||
else:
|
||||
return self.get_project_db(file_path)
|
||||
|
||||
def close_all(self):
|
||||
"""Close all project databases."""
|
||||
for db in self.projects.values():
|
||||
db.close()
|
||||
self.projects.clear()
|
||||
@@ -1,528 +0,0 @@
|
||||
"""Unified graph database tool."""
|
||||
|
||||
import json
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
Unpack,
|
||||
Optional,
|
||||
Annotated,
|
||||
TypedDict,
|
||||
final,
|
||||
override,
|
||||
)
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.database.database_manager import DatabaseManager
|
||||
|
||||
# Parameter types
|
||||
Action = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Action: query (default), add, remove, search, stats",
|
||||
default="query",
|
||||
),
|
||||
]
|
||||
|
||||
NodeId = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Node ID",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
NodeType = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Node type/label",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
EdgeType = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Edge type/relationship",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
FromNode = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Source node ID for edges",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
ToNode = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Target node ID for edges",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Properties = Annotated[
|
||||
Optional[Dict[str, Any]],
|
||||
Field(
|
||||
description="Node/edge properties as JSON",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Pattern = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Search pattern for properties",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Depth = Annotated[
|
||||
int,
|
||||
Field(
|
||||
description="Max traversal depth for queries",
|
||||
default=2,
|
||||
),
|
||||
]
|
||||
|
||||
Limit = Annotated[
|
||||
int,
|
||||
Field(
|
||||
description="Maximum results to return",
|
||||
default=50,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class GraphParams(TypedDict, total=False):
|
||||
"""Parameters for graph tool."""
|
||||
|
||||
action: str
|
||||
node_id: Optional[str]
|
||||
node_type: Optional[str]
|
||||
edge_type: Optional[str]
|
||||
from_node: Optional[str]
|
||||
to_node: Optional[str]
|
||||
properties: Optional[Dict[str, Any]]
|
||||
pattern: Optional[str]
|
||||
depth: int
|
||||
limit: int
|
||||
|
||||
|
||||
@final
|
||||
class GraphTool(BaseTool):
|
||||
"""Unified graph database tool."""
|
||||
|
||||
def __init__(
|
||||
self, permission_manager: PermissionManager, db_manager: DatabaseManager
|
||||
):
|
||||
"""Initialize the graph tool."""
|
||||
super().__init__(permission_manager)
|
||||
self.db_manager = db_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "graph"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Graph database. Actions: query (default), add, remove, search, stats.
|
||||
|
||||
Usage:
|
||||
graph --node-id user123
|
||||
graph --action add --node-id user123 --node-type User --properties '{"name": "John"}'
|
||||
graph --action add --from-node user123 --to-node post456 --edge-type CREATED
|
||||
graph --action query --node-id user123 --depth 3
|
||||
graph --action search --pattern "John" --node-type User
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[GraphParams],
|
||||
) -> str:
|
||||
"""Execute graph operation."""
|
||||
tool_ctx = self.create_tool_context(ctx)
|
||||
|
||||
# Get current project database
|
||||
project_db = self.db_manager.get_current_project_db()
|
||||
if not project_db:
|
||||
return "Error: No project database found. Are you in a project directory?"
|
||||
|
||||
# Extract action
|
||||
action = params.get("action", "query")
|
||||
|
||||
# Route to appropriate handler
|
||||
if action == "query":
|
||||
return await self._handle_query(project_db, params, tool_ctx)
|
||||
elif action == "add":
|
||||
return await self._handle_add(project_db, params, tool_ctx)
|
||||
elif action == "remove":
|
||||
return await self._handle_remove(project_db, params, tool_ctx)
|
||||
elif action == "search":
|
||||
return await self._handle_search(project_db, params, tool_ctx)
|
||||
elif action == "stats":
|
||||
return await self._handle_stats(project_db, tool_ctx)
|
||||
else:
|
||||
return f"Error: Unknown action '{action}'. Valid actions: query, add, remove, search, stats"
|
||||
|
||||
async def _handle_query(self, project_db, params: Dict[str, Any], tool_ctx) -> str:
|
||||
"""Query graph relationships."""
|
||||
node_id = params.get("node_id")
|
||||
node_type = params.get("node_type")
|
||||
depth = params.get("depth", 2)
|
||||
limit = params.get("limit", 50)
|
||||
|
||||
if not node_id and not node_type:
|
||||
return "Error: node_id or node_type required for query"
|
||||
|
||||
try:
|
||||
with project_db.get_graph_connection() as conn:
|
||||
results = []
|
||||
|
||||
if node_id:
|
||||
# Query specific node and its relationships
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
WITH RECURSIVE
|
||||
node_tree(id, type, properties, depth, path) AS (
|
||||
SELECT id, type, properties, 0, id
|
||||
FROM nodes
|
||||
WHERE id = ?
|
||||
|
||||
UNION ALL
|
||||
|
||||
SELECT n.id, n.type, n.properties, nt.depth + 1,
|
||||
nt.path || ' -> ' || n.id
|
||||
FROM nodes n
|
||||
JOIN edges e ON (e.to_node = n.id OR e.from_node = n.id)
|
||||
JOIN node_tree nt ON (
|
||||
(e.from_node = nt.id AND e.to_node = n.id) OR
|
||||
(e.to_node = nt.id AND e.from_node = n.id)
|
||||
)
|
||||
WHERE nt.depth < ?
|
||||
)
|
||||
SELECT DISTINCT * FROM node_tree
|
||||
ORDER BY depth, id
|
||||
LIMIT ?
|
||||
""",
|
||||
(node_id, depth, limit),
|
||||
)
|
||||
|
||||
nodes = cursor.fetchall()
|
||||
|
||||
# Get edges
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
SELECT from_node, to_node, type, properties
|
||||
FROM edges
|
||||
WHERE from_node IN (SELECT id FROM node_tree)
|
||||
OR to_node IN (SELECT id FROM node_tree)
|
||||
"""
|
||||
)
|
||||
|
||||
edges = cursor.fetchall()
|
||||
|
||||
else:
|
||||
# Query by type
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
SELECT id, type, properties
|
||||
FROM nodes
|
||||
WHERE type = ?
|
||||
LIMIT ?
|
||||
""",
|
||||
(node_type, limit),
|
||||
)
|
||||
|
||||
nodes = cursor.fetchall()
|
||||
edges = []
|
||||
|
||||
# Format results
|
||||
output = ["=== Graph Query Results ==="]
|
||||
|
||||
if nodes:
|
||||
output.append(f"\nNodes ({len(nodes)}):")
|
||||
for node in nodes:
|
||||
props = json.loads(node[2]) if node[2] else {}
|
||||
output.append(f" {node[0]} [{node[1]}] {props}")
|
||||
|
||||
if edges:
|
||||
output.append(f"\nEdges ({len(edges)}):")
|
||||
for edge in edges:
|
||||
props = json.loads(edge[3]) if edge[3] else {}
|
||||
output.append(f" {edge[0]} --[{edge[2]}]--> {edge[1]} {props}")
|
||||
|
||||
if not nodes and not edges:
|
||||
output.append("No results found")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Query failed: {str(e)}")
|
||||
return f"Error during query: {str(e)}"
|
||||
|
||||
async def _handle_add(self, project_db, params: Dict[str, Any], tool_ctx) -> str:
|
||||
"""Add nodes or edges."""
|
||||
node_id = params.get("node_id")
|
||||
from_node = params.get("from_node")
|
||||
to_node = params.get("to_node")
|
||||
|
||||
if node_id:
|
||||
# Add node
|
||||
node_type = params.get("node_type")
|
||||
if not node_type:
|
||||
return "Error: node_type required when adding node"
|
||||
|
||||
properties = params.get("properties", {})
|
||||
|
||||
try:
|
||||
with project_db.get_graph_connection() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT OR REPLACE INTO nodes (id, type, properties)
|
||||
VALUES (?, ?, ?)
|
||||
""",
|
||||
(node_id, node_type, json.dumps(properties)),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
await tool_ctx.info(f"Added node: {node_id}")
|
||||
return f"Added node {node_id} [{node_type}]"
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to add node: {str(e)}")
|
||||
return f"Error adding node: {str(e)}"
|
||||
|
||||
elif from_node and to_node:
|
||||
# Add edge
|
||||
edge_type = params.get("edge_type", "RELATED")
|
||||
properties = params.get("properties", {})
|
||||
|
||||
try:
|
||||
with project_db.get_graph_connection() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT OR REPLACE INTO edges (from_node, to_node, type, properties)
|
||||
VALUES (?, ?, ?, ?)
|
||||
""",
|
||||
(from_node, to_node, edge_type, json.dumps(properties)),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
await tool_ctx.info(f"Added edge: {from_node} -> {to_node}")
|
||||
return f"Added edge {from_node} --[{edge_type}]--> {to_node}"
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to add edge: {str(e)}")
|
||||
return f"Error adding edge: {str(e)}"
|
||||
|
||||
else:
|
||||
return "Error: Either node_id (for node) or from_node + to_node (for edge) required"
|
||||
|
||||
async def _handle_remove(self, project_db, params: Dict[str, Any], tool_ctx) -> str:
|
||||
"""Remove nodes or edges."""
|
||||
node_id = params.get("node_id")
|
||||
from_node = params.get("from_node")
|
||||
to_node = params.get("to_node")
|
||||
|
||||
if node_id:
|
||||
# Remove node and its edges
|
||||
try:
|
||||
with project_db.get_graph_connection() as conn:
|
||||
# Delete edges first
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
DELETE FROM edges
|
||||
WHERE from_node = ? OR to_node = ?
|
||||
""",
|
||||
(node_id, node_id),
|
||||
)
|
||||
|
||||
edges_deleted = cursor.rowcount
|
||||
|
||||
# Delete node
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
DELETE FROM nodes WHERE id = ?
|
||||
""",
|
||||
(node_id,),
|
||||
)
|
||||
|
||||
if cursor.rowcount == 0:
|
||||
return f"Node {node_id} not found"
|
||||
|
||||
conn.commit()
|
||||
|
||||
msg = f"Removed node {node_id}"
|
||||
if edges_deleted > 0:
|
||||
msg += f" and {edges_deleted} connected edges"
|
||||
|
||||
await tool_ctx.info(msg)
|
||||
return msg
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to remove node: {str(e)}")
|
||||
return f"Error removing node: {str(e)}"
|
||||
|
||||
elif from_node and to_node:
|
||||
# Remove specific edge
|
||||
edge_type = params.get("edge_type")
|
||||
|
||||
try:
|
||||
with project_db.get_graph_connection() as conn:
|
||||
if edge_type:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
DELETE FROM edges
|
||||
WHERE from_node = ? AND to_node = ? AND type = ?
|
||||
""",
|
||||
(from_node, to_node, edge_type),
|
||||
)
|
||||
else:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
DELETE FROM edges
|
||||
WHERE from_node = ? AND to_node = ?
|
||||
""",
|
||||
(from_node, to_node),
|
||||
)
|
||||
|
||||
if cursor.rowcount == 0:
|
||||
return f"Edge not found"
|
||||
|
||||
conn.commit()
|
||||
|
||||
return f"Removed edge {from_node} --> {to_node}"
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to remove edge: {str(e)}")
|
||||
return f"Error removing edge: {str(e)}"
|
||||
|
||||
else:
|
||||
return "Error: Either node_id or from_node + to_node required for remove"
|
||||
|
||||
async def _handle_search(self, project_db, params: Dict[str, Any], tool_ctx) -> str:
|
||||
"""Search graph by pattern."""
|
||||
pattern = params.get("pattern")
|
||||
if not pattern:
|
||||
return "Error: pattern required for search"
|
||||
|
||||
node_type = params.get("node_type")
|
||||
limit = params.get("limit", 50)
|
||||
|
||||
try:
|
||||
with project_db.get_graph_connection() as conn:
|
||||
# Search in properties
|
||||
if node_type:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
SELECT id, type, properties
|
||||
FROM nodes
|
||||
WHERE type = ? AND properties LIKE ?
|
||||
LIMIT ?
|
||||
""",
|
||||
(node_type, f"%{pattern}%", limit),
|
||||
)
|
||||
else:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
SELECT id, type, properties
|
||||
FROM nodes
|
||||
WHERE properties LIKE ?
|
||||
LIMIT ?
|
||||
""",
|
||||
(f"%{pattern}%", limit),
|
||||
)
|
||||
|
||||
results = cursor.fetchall()
|
||||
|
||||
if not results:
|
||||
return f"No nodes found matching '{pattern}'"
|
||||
|
||||
# Format results
|
||||
output = [f"=== Graph Search Results for '{pattern}' ==="]
|
||||
output.append(f"Found {len(results)} nodes\n")
|
||||
|
||||
for node in results:
|
||||
props = json.loads(node[2]) if node[2] else {}
|
||||
output.append(f"{node[0]} [{node[1]}] {props}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Search failed: {str(e)}")
|
||||
return f"Error during search: {str(e)}"
|
||||
|
||||
async def _handle_stats(self, project_db, tool_ctx) -> str:
|
||||
"""Get graph statistics."""
|
||||
try:
|
||||
with project_db.get_graph_connection() as conn:
|
||||
# Node stats
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
SELECT type, COUNT(*) as count
|
||||
FROM nodes
|
||||
GROUP BY type
|
||||
ORDER BY count DESC
|
||||
"""
|
||||
)
|
||||
|
||||
node_stats = cursor.fetchall()
|
||||
|
||||
# Edge stats
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
SELECT type, COUNT(*) as count
|
||||
FROM edges
|
||||
GROUP BY type
|
||||
ORDER BY count DESC
|
||||
"""
|
||||
)
|
||||
|
||||
edge_stats = cursor.fetchall()
|
||||
|
||||
# Total counts
|
||||
cursor = conn.execute("SELECT COUNT(*) FROM nodes")
|
||||
total_nodes = cursor.fetchone()[0]
|
||||
|
||||
cursor = conn.execute("SELECT COUNT(*) FROM edges")
|
||||
total_edges = cursor.fetchone()[0]
|
||||
|
||||
# Format output
|
||||
output = [f"=== Graph Database Statistics ==="]
|
||||
output.append(f"Project: {project_db.project_path}")
|
||||
output.append(f"\nTotal nodes: {total_nodes}")
|
||||
output.append(f"Total edges: {total_edges}")
|
||||
|
||||
if node_stats:
|
||||
output.append("\nNodes by type:")
|
||||
for node_type, count in node_stats:
|
||||
output.append(f" {node_type}: {count}")
|
||||
|
||||
if edge_stats:
|
||||
output.append("\nEdges by type:")
|
||||
for edge_type, count in edge_stats:
|
||||
output.append(f" {edge_type}: {count}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to get stats: {str(e)}")
|
||||
return f"Error getting stats: {str(e)}"
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,269 +0,0 @@
|
||||
"""Graph add tool for adding nodes and edges to the graph database."""
|
||||
|
||||
import json
|
||||
from typing import Unpack, Optional, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.database.database_manager import DatabaseManager
|
||||
|
||||
NodeId = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Node ID to add (required for nodes)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
NodeType = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Node type (e.g., 'file', 'function', 'class')",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Properties = Annotated[
|
||||
Optional[dict],
|
||||
Field(
|
||||
description="Properties as JSON object",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Source = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Source node ID (required for edges)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Target = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Target node ID (required for edges)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Relationship = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Edge relationship type (e.g., 'imports', 'calls', 'inherits')",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Weight = Annotated[
|
||||
float,
|
||||
Field(
|
||||
description="Edge weight (default 1.0)",
|
||||
default=1.0,
|
||||
),
|
||||
]
|
||||
|
||||
ProjectPath = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Project path (defaults to current directory)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class GraphAddParams(TypedDict, total=False):
|
||||
"""Parameters for graph add tool."""
|
||||
|
||||
node_id: Optional[str]
|
||||
node_type: Optional[str]
|
||||
properties: Optional[dict]
|
||||
source: Optional[str]
|
||||
target: Optional[str]
|
||||
relationship: Optional[str]
|
||||
weight: float
|
||||
project_path: Optional[str]
|
||||
|
||||
|
||||
@final
|
||||
class GraphAddTool(BaseTool):
|
||||
"""Tool for adding nodes and edges to graph database."""
|
||||
|
||||
def __init__(
|
||||
self, permission_manager: PermissionManager, db_manager: DatabaseManager
|
||||
):
|
||||
"""Initialize the graph add tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
db_manager: Database manager instance
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.db_manager = db_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "graph_add"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Add nodes and edges to the project's graph database.
|
||||
|
||||
To add a node:
|
||||
- Provide node_id and node_type
|
||||
- Optionally add properties as JSON
|
||||
|
||||
To add an edge:
|
||||
- Provide source, target, and relationship
|
||||
- Optionally add weight and properties
|
||||
|
||||
Common node types:
|
||||
- file, function, class, module, variable
|
||||
|
||||
Common relationships:
|
||||
- imports, calls, inherits, contains, references, depends_on
|
||||
|
||||
Examples:
|
||||
- graph_add --node-id "main.py" --node-type "file" --properties '{"size": 1024}'
|
||||
- graph_add --node-id "MyClass" --node-type "class" --properties '{"file": "main.py"}'
|
||||
- graph_add --source "main.py" --target "utils.py" --relationship "imports"
|
||||
- graph_add --source "func1" --target "func2" --relationship "calls" --weight 5.0
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[GraphAddParams],
|
||||
) -> str:
|
||||
"""Add nodes or edges to graph.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Result of add operation
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
node_id = params.get("node_id")
|
||||
node_type = params.get("node_type")
|
||||
properties = params.get("properties", {})
|
||||
source = params.get("source")
|
||||
target = params.get("target")
|
||||
relationship = params.get("relationship")
|
||||
weight = params.get("weight", 1.0)
|
||||
project_path = params.get("project_path")
|
||||
|
||||
# Determine if adding node or edge
|
||||
is_node = node_id is not None
|
||||
is_edge = source is not None and target is not None
|
||||
|
||||
if not is_node and not is_edge:
|
||||
return "Error: Must provide either (node_id and node_type) for a node, or (source, target, relationship) for an edge"
|
||||
|
||||
if is_node and is_edge:
|
||||
return "Error: Cannot add both node and edge in one operation"
|
||||
|
||||
# Get project database
|
||||
try:
|
||||
if project_path:
|
||||
project_db = self.db_manager.get_project_db(project_path)
|
||||
else:
|
||||
import os
|
||||
|
||||
project_db = self.db_manager.get_project_for_path(os.getcwd())
|
||||
|
||||
if not project_db:
|
||||
return "Error: Could not find project database"
|
||||
|
||||
except PermissionError as e:
|
||||
return str(e)
|
||||
except Exception as e:
|
||||
return f"Error accessing project database: {str(e)}"
|
||||
|
||||
# Get graph connection
|
||||
graph_conn = project_db.get_graph_connection()
|
||||
|
||||
try:
|
||||
if is_node:
|
||||
# Add node
|
||||
if not node_type:
|
||||
return "Error: node_type is required when adding a node"
|
||||
|
||||
await tool_ctx.info(f"Adding node: {node_id} (type: {node_type})")
|
||||
|
||||
# Serialize properties
|
||||
properties_json = json.dumps(properties) if properties else None
|
||||
|
||||
# Insert or update node
|
||||
graph_conn.execute(
|
||||
"""
|
||||
INSERT OR REPLACE INTO nodes (id, type, properties)
|
||||
VALUES (?, ?, ?)
|
||||
""",
|
||||
(node_id, node_type, properties_json),
|
||||
)
|
||||
|
||||
graph_conn.commit()
|
||||
|
||||
return f"Successfully added node '{node_id}' of type '{node_type}'"
|
||||
|
||||
else:
|
||||
# Add edge
|
||||
if not relationship:
|
||||
return "Error: relationship is required when adding an edge"
|
||||
|
||||
await tool_ctx.info(
|
||||
f"Adding edge: {source} --[{relationship}]--> {target}"
|
||||
)
|
||||
|
||||
# Check if nodes exist
|
||||
cursor = graph_conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT id FROM nodes WHERE id IN (?, ?)", (source, target)
|
||||
)
|
||||
existing = [row[0] for row in cursor.fetchall()]
|
||||
|
||||
if source not in existing:
|
||||
return f"Error: Source node '{source}' does not exist"
|
||||
if target not in existing:
|
||||
return f"Error: Target node '{target}' does not exist"
|
||||
|
||||
# Serialize properties
|
||||
properties_json = json.dumps(properties) if properties else None
|
||||
|
||||
# Insert or update edge
|
||||
graph_conn.execute(
|
||||
"""
|
||||
INSERT OR REPLACE INTO edges (source, target, relationship, weight, properties)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
""",
|
||||
(source, target, relationship, weight, properties_json),
|
||||
)
|
||||
|
||||
graph_conn.commit()
|
||||
|
||||
# Save to disk
|
||||
project_db._save_graph_to_disk()
|
||||
|
||||
return f"Successfully added edge: {source} --[{relationship}]--> {target} (weight: {weight})"
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to add to graph: {str(e)}")
|
||||
return f"Error adding to graph: {str(e)}"
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,612 +0,0 @@
|
||||
"""Graph query tool for querying the graph database."""
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
from typing import (
|
||||
Unpack,
|
||||
Optional,
|
||||
Annotated,
|
||||
TypedDict,
|
||||
final,
|
||||
override,
|
||||
)
|
||||
from collections import deque
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.database.database_manager import DatabaseManager
|
||||
|
||||
Query = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Query type: neighbors, path, subgraph, connected, ancestors, descendants",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
NodeId = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Starting node ID",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
TargetId = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Target node ID (for path queries)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Depth = Annotated[
|
||||
int,
|
||||
Field(
|
||||
description="Maximum depth for traversal",
|
||||
default=2,
|
||||
),
|
||||
]
|
||||
|
||||
Relationship = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Filter by relationship type",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
NodeType = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Filter by node type",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Direction = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Direction: both, incoming, outgoing",
|
||||
default="both",
|
||||
),
|
||||
]
|
||||
|
||||
ProjectPath = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Project path (defaults to current directory)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class GraphQueryParams(TypedDict, total=False):
|
||||
"""Parameters for graph query tool."""
|
||||
|
||||
query: str
|
||||
node_id: Optional[str]
|
||||
target_id: Optional[str]
|
||||
depth: int
|
||||
relationship: Optional[str]
|
||||
node_type: Optional[str]
|
||||
direction: str
|
||||
project_path: Optional[str]
|
||||
|
||||
|
||||
@final
|
||||
class GraphQueryTool(BaseTool):
|
||||
"""Tool for querying the graph database."""
|
||||
|
||||
def __init__(
|
||||
self, permission_manager: PermissionManager, db_manager: DatabaseManager
|
||||
):
|
||||
"""Initialize the graph query tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
db_manager: Database manager instance
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.db_manager = db_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "graph_query"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Query the project's graph database for relationships and patterns.
|
||||
|
||||
Query types:
|
||||
- neighbors: Find direct neighbors of a node
|
||||
- path: Find shortest path between two nodes
|
||||
- subgraph: Get subgraph around a node up to depth
|
||||
- connected: Find all nodes connected to a node
|
||||
- ancestors: Find nodes that point TO this node
|
||||
- descendants: Find nodes that this node points TO
|
||||
|
||||
Options:
|
||||
- --depth: Max traversal depth (default 2)
|
||||
- --relationship: Filter by edge type
|
||||
- --node-type: Filter by node type
|
||||
- --direction: both, incoming, outgoing
|
||||
|
||||
Examples:
|
||||
- graph_query --query neighbors --node-id "main.py"
|
||||
- graph_query --query path --node-id "main.py" --target-id "utils.py"
|
||||
- graph_query --query subgraph --node-id "MyClass" --depth 3
|
||||
- graph_query --query ancestors --node-id "error_handler" --relationship "calls"
|
||||
- graph_query --query descendants --node-id "BaseClass" --relationship "inherits"
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[GraphQueryParams],
|
||||
) -> str:
|
||||
"""Execute graph query.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Query results
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
query = params.get("query")
|
||||
if not query:
|
||||
return "Error: query is required"
|
||||
|
||||
node_id = params.get("node_id")
|
||||
target_id = params.get("target_id")
|
||||
depth = params.get("depth", 2)
|
||||
relationship = params.get("relationship")
|
||||
node_type = params.get("node_type")
|
||||
direction = params.get("direction", "both")
|
||||
project_path = params.get("project_path")
|
||||
|
||||
# Validate query type
|
||||
valid_queries = [
|
||||
"neighbors",
|
||||
"path",
|
||||
"subgraph",
|
||||
"connected",
|
||||
"ancestors",
|
||||
"descendants",
|
||||
]
|
||||
if query not in valid_queries:
|
||||
return f"Error: Invalid query '{query}'. Must be one of: {', '.join(valid_queries)}"
|
||||
|
||||
# Validate required parameters
|
||||
if (
|
||||
query in ["neighbors", "subgraph", "connected", "ancestors", "descendants"]
|
||||
and not node_id
|
||||
):
|
||||
return f"Error: node_id is required for '{query}' query"
|
||||
|
||||
if query == "path" and (not node_id or not target_id):
|
||||
return "Error: Both node_id and target_id are required for 'path' query"
|
||||
|
||||
# Get project database
|
||||
try:
|
||||
if project_path:
|
||||
project_db = self.db_manager.get_project_db(project_path)
|
||||
else:
|
||||
import os
|
||||
|
||||
project_db = self.db_manager.get_project_for_path(os.getcwd())
|
||||
|
||||
if not project_db:
|
||||
return "Error: Could not find project database"
|
||||
|
||||
except PermissionError as e:
|
||||
return str(e)
|
||||
except Exception as e:
|
||||
return f"Error accessing project database: {str(e)}"
|
||||
|
||||
# Get graph connection
|
||||
graph_conn = project_db.get_graph_connection()
|
||||
|
||||
await tool_ctx.info(f"Executing {query} query")
|
||||
|
||||
try:
|
||||
if query == "neighbors":
|
||||
return self._query_neighbors(
|
||||
graph_conn, node_id, relationship, node_type, direction
|
||||
)
|
||||
elif query == "path":
|
||||
return self._query_path(graph_conn, node_id, target_id, relationship)
|
||||
elif query == "subgraph":
|
||||
return self._query_subgraph(
|
||||
graph_conn, node_id, depth, relationship, node_type, direction
|
||||
)
|
||||
elif query == "connected":
|
||||
return self._query_connected(
|
||||
graph_conn, node_id, relationship, node_type, direction
|
||||
)
|
||||
elif query == "ancestors":
|
||||
return self._query_ancestors(
|
||||
graph_conn, node_id, depth, relationship, node_type
|
||||
)
|
||||
elif query == "descendants":
|
||||
return self._query_descendants(
|
||||
graph_conn, node_id, depth, relationship, node_type
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to execute query: {str(e)}")
|
||||
return f"Error executing query: {str(e)}"
|
||||
|
||||
def _query_neighbors(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
node_id: str,
|
||||
relationship: Optional[str],
|
||||
node_type: Optional[str],
|
||||
direction: str,
|
||||
) -> str:
|
||||
"""Get direct neighbors of a node."""
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Check if node exists
|
||||
cursor.execute("SELECT type, properties FROM nodes WHERE id = ?", (node_id,))
|
||||
node_info = cursor.fetchone()
|
||||
if not node_info:
|
||||
return f"Error: Node '{node_id}' not found"
|
||||
|
||||
neighbors = []
|
||||
|
||||
# Get outgoing edges
|
||||
if direction in ["both", "outgoing"]:
|
||||
query = """SELECT e.target, e.relationship, e.weight, n.type, n.properties
|
||||
FROM edges e JOIN nodes n ON e.target = n.id
|
||||
WHERE e.source = ?"""
|
||||
params = [node_id]
|
||||
|
||||
if relationship:
|
||||
query += " AND e.relationship = ?"
|
||||
params.append(relationship)
|
||||
if node_type:
|
||||
query += " AND n.type = ?"
|
||||
params.append(node_type)
|
||||
|
||||
cursor.execute(query, params)
|
||||
for row in cursor.fetchall():
|
||||
neighbors.append(
|
||||
{
|
||||
"direction": "outgoing",
|
||||
"node_id": row[0],
|
||||
"relationship": row[1],
|
||||
"weight": row[2],
|
||||
"node_type": row[3],
|
||||
"properties": json.loads(row[4]) if row[4] else {},
|
||||
}
|
||||
)
|
||||
|
||||
# Get incoming edges
|
||||
if direction in ["both", "incoming"]:
|
||||
query = """SELECT e.source, e.relationship, e.weight, n.type, n.properties
|
||||
FROM edges e JOIN nodes n ON e.source = n.id
|
||||
WHERE e.target = ?"""
|
||||
params = [node_id]
|
||||
|
||||
if relationship:
|
||||
query += " AND e.relationship = ?"
|
||||
params.append(relationship)
|
||||
if node_type:
|
||||
query += " AND n.type = ?"
|
||||
params.append(node_type)
|
||||
|
||||
cursor.execute(query, params)
|
||||
for row in cursor.fetchall():
|
||||
neighbors.append(
|
||||
{
|
||||
"direction": "incoming",
|
||||
"node_id": row[0],
|
||||
"relationship": row[1],
|
||||
"weight": row[2],
|
||||
"node_type": row[3],
|
||||
"properties": json.loads(row[4]) if row[4] else {},
|
||||
}
|
||||
)
|
||||
|
||||
if not neighbors:
|
||||
return f"No neighbors found for node '{node_id}'"
|
||||
|
||||
# Format output
|
||||
output = [f"Neighbors of '{node_id}' ({node_info[0]}):\n"]
|
||||
for n in neighbors:
|
||||
arrow = "<--" if n["direction"] == "incoming" else "-->"
|
||||
output.append(
|
||||
f" {node_id} {arrow}[{n['relationship']}]--> {n['node_id']} ({n['node_type']})"
|
||||
)
|
||||
if n["properties"]:
|
||||
output.append(
|
||||
f" Properties: {json.dumps(n['properties'], indent=6)[:100]}"
|
||||
)
|
||||
|
||||
output.append(f"\nTotal neighbors: {len(neighbors)}")
|
||||
return "\n".join(output)
|
||||
|
||||
def _query_path(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
start: str,
|
||||
end: str,
|
||||
relationship: Optional[str],
|
||||
) -> str:
|
||||
"""Find shortest path between two nodes using BFS."""
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Check if nodes exist
|
||||
cursor.execute("SELECT id FROM nodes WHERE id IN (?, ?)", (start, end))
|
||||
existing = [row[0] for row in cursor.fetchall()]
|
||||
if start not in existing:
|
||||
return f"Error: Start node '{start}' not found"
|
||||
if end not in existing:
|
||||
return f"Error: End node '{end}' not found"
|
||||
|
||||
# BFS to find shortest path
|
||||
queue = deque([(start, [start])])
|
||||
visited = {start}
|
||||
|
||||
while queue:
|
||||
current, path = queue.popleft()
|
||||
|
||||
if current == end:
|
||||
# Found path, get edge details
|
||||
output = [f"Shortest path from '{start}' to '{end}':\n"]
|
||||
|
||||
for i in range(len(path) - 1):
|
||||
src, tgt = path[i], path[i + 1]
|
||||
|
||||
# Get edge details
|
||||
query = "SELECT relationship, weight FROM edges WHERE source = ? AND target = ?"
|
||||
cursor.execute(query, (src, tgt))
|
||||
edge = cursor.fetchone()
|
||||
|
||||
if edge:
|
||||
output.append(f" {src} --[{edge[0]}]--> {tgt}")
|
||||
else:
|
||||
output.append(f" {src} --> {tgt}")
|
||||
|
||||
output.append(f"\nPath length: {len(path) - 1} edge(s)")
|
||||
return "\n".join(output)
|
||||
|
||||
# Get neighbors
|
||||
query = "SELECT target FROM edges WHERE source = ?"
|
||||
params = [current]
|
||||
if relationship:
|
||||
query += " AND relationship = ?"
|
||||
params.append(relationship)
|
||||
|
||||
cursor.execute(query, params)
|
||||
|
||||
for (neighbor,) in cursor.fetchall():
|
||||
if neighbor not in visited:
|
||||
visited.add(neighbor)
|
||||
queue.append((neighbor, path + [neighbor]))
|
||||
|
||||
return f"No path found from '{start}' to '{end}'" + (
|
||||
f" with relationship '{relationship}'" if relationship else ""
|
||||
)
|
||||
|
||||
def _query_subgraph(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
node_id: str,
|
||||
depth: int,
|
||||
relationship: Optional[str],
|
||||
node_type: Optional[str],
|
||||
direction: str,
|
||||
) -> str:
|
||||
"""Get subgraph around a node up to specified depth."""
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Check if node exists
|
||||
cursor.execute("SELECT type FROM nodes WHERE id = ?", (node_id,))
|
||||
if not cursor.fetchone():
|
||||
return f"Error: Node '{node_id}' not found"
|
||||
|
||||
# BFS to collect nodes and edges
|
||||
nodes = {node_id: 0} # node_id -> depth
|
||||
edges = set() # (source, target, relationship)
|
||||
queue = deque([(node_id, 0)])
|
||||
|
||||
while queue:
|
||||
current, current_depth = queue.popleft()
|
||||
|
||||
if current_depth >= depth:
|
||||
continue
|
||||
|
||||
# Get edges based on direction
|
||||
if direction in ["both", "outgoing"]:
|
||||
query = """SELECT e.target, e.relationship, n.type
|
||||
FROM edges e JOIN nodes n ON e.target = n.id
|
||||
WHERE e.source = ?"""
|
||||
params = [current]
|
||||
|
||||
if relationship:
|
||||
query += " AND e.relationship = ?"
|
||||
params.append(relationship)
|
||||
if node_type:
|
||||
query += " AND n.type = ?"
|
||||
params.append(node_type)
|
||||
|
||||
cursor.execute(query, params)
|
||||
|
||||
for target, rel, _ in cursor.fetchall():
|
||||
edges.add((current, target, rel))
|
||||
if target not in nodes or nodes[target] > current_depth + 1:
|
||||
nodes[target] = current_depth + 1
|
||||
queue.append((target, current_depth + 1))
|
||||
|
||||
if direction in ["both", "incoming"]:
|
||||
query = """SELECT e.source, e.relationship, n.type
|
||||
FROM edges e JOIN nodes n ON e.source = n.id
|
||||
WHERE e.target = ?"""
|
||||
params = [current]
|
||||
|
||||
if relationship:
|
||||
query += " AND e.relationship = ?"
|
||||
params.append(relationship)
|
||||
if node_type:
|
||||
query += " AND n.type = ?"
|
||||
params.append(node_type)
|
||||
|
||||
cursor.execute(query, params)
|
||||
|
||||
for source, rel, _ in cursor.fetchall():
|
||||
edges.add((source, current, rel))
|
||||
if source not in nodes or nodes[source] > current_depth + 1:
|
||||
nodes[source] = current_depth + 1
|
||||
queue.append((source, current_depth + 1))
|
||||
|
||||
# Format output
|
||||
output = [f"Subgraph around '{node_id}' (depth={depth}):\n"]
|
||||
output.append(f"Nodes ({len(nodes)}):")
|
||||
|
||||
# Get node details
|
||||
for node, d in sorted(nodes.items(), key=lambda x: (x[1], x[0])):
|
||||
cursor.execute("SELECT type FROM nodes WHERE id = ?", (node,))
|
||||
node_type = cursor.fetchone()[0]
|
||||
output.append(f" [{d}] {node} ({node_type})")
|
||||
|
||||
output.append(f"\nEdges ({len(edges)}):")
|
||||
for src, tgt, rel in sorted(edges):
|
||||
output.append(f" {src} --[{rel}]--> {tgt}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _query_connected(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
node_id: str,
|
||||
relationship: Optional[str],
|
||||
node_type: Optional[str],
|
||||
direction: str,
|
||||
) -> str:
|
||||
"""Find all nodes connected to a node (transitive closure)."""
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Check if node exists
|
||||
cursor.execute("SELECT type FROM nodes WHERE id = ?", (node_id,))
|
||||
if not cursor.fetchone():
|
||||
return f"Error: Node '{node_id}' not found"
|
||||
|
||||
# BFS to find all connected nodes
|
||||
visited = {node_id}
|
||||
queue = deque([node_id])
|
||||
connections = [] # (node_id, node_type, distance)
|
||||
distance = {node_id: 0}
|
||||
|
||||
while queue:
|
||||
current = queue.popleft()
|
||||
current_dist = distance[current]
|
||||
|
||||
# Get edges based on direction
|
||||
neighbors = []
|
||||
|
||||
if direction in ["both", "outgoing"]:
|
||||
query = """SELECT e.target, n.type FROM edges e
|
||||
JOIN nodes n ON e.target = n.id
|
||||
WHERE e.source = ?"""
|
||||
params = [current]
|
||||
|
||||
if relationship:
|
||||
query += " AND e.relationship = ?"
|
||||
params.append(relationship)
|
||||
if node_type:
|
||||
query += " AND n.type = ?"
|
||||
params.append(node_type)
|
||||
|
||||
cursor.execute(query, params)
|
||||
neighbors.extend(cursor.fetchall())
|
||||
|
||||
if direction in ["both", "incoming"]:
|
||||
query = """SELECT e.source, n.type FROM edges e
|
||||
JOIN nodes n ON e.source = n.id
|
||||
WHERE e.target = ?"""
|
||||
params = [current]
|
||||
|
||||
if relationship:
|
||||
query += " AND e.relationship = ?"
|
||||
params.append(relationship)
|
||||
if node_type:
|
||||
query += " AND n.type = ?"
|
||||
params.append(node_type)
|
||||
|
||||
cursor.execute(query, params)
|
||||
neighbors.extend(cursor.fetchall())
|
||||
|
||||
for neighbor, n_type in neighbors:
|
||||
if neighbor not in visited:
|
||||
visited.add(neighbor)
|
||||
queue.append(neighbor)
|
||||
distance[neighbor] = current_dist + 1
|
||||
connections.append((neighbor, n_type, current_dist + 1))
|
||||
|
||||
if not connections:
|
||||
return f"No connected nodes found for '{node_id}'"
|
||||
|
||||
# Format output
|
||||
output = [f"Nodes connected to '{node_id}' ({direction}):"]
|
||||
output.append(f"\nTotal connected: {len(connections)}\n")
|
||||
|
||||
# Group by distance
|
||||
by_distance = {}
|
||||
for node, n_type, dist in connections:
|
||||
if dist not in by_distance:
|
||||
by_distance[dist] = []
|
||||
by_distance[dist].append((node, n_type))
|
||||
|
||||
for dist in sorted(by_distance.keys()):
|
||||
output.append(f"Distance {dist}:")
|
||||
for node, n_type in sorted(by_distance[dist]):
|
||||
output.append(f" {node} ({n_type})")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _query_ancestors(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
node_id: str,
|
||||
depth: int,
|
||||
relationship: Optional[str],
|
||||
node_type: Optional[str],
|
||||
) -> str:
|
||||
"""Find nodes that point TO this node (incoming edges only)."""
|
||||
return self._query_subgraph(
|
||||
conn, node_id, depth, relationship, node_type, "incoming"
|
||||
)
|
||||
|
||||
def _query_descendants(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
node_id: str,
|
||||
depth: int,
|
||||
relationship: Optional[str],
|
||||
node_type: Optional[str],
|
||||
) -> str:
|
||||
"""Find nodes that this node points TO (outgoing edges only)."""
|
||||
return self._query_subgraph(
|
||||
conn, node_id, depth, relationship, node_type, "outgoing"
|
||||
)
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,272 +0,0 @@
|
||||
"""Graph remove tool for removing nodes and edges from the graph database."""
|
||||
|
||||
from typing import Unpack, Optional, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.database.database_manager import DatabaseManager
|
||||
|
||||
NodeId = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Node ID to remove",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Source = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Source node ID (for edge removal)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Target = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Target node ID (for edge removal)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Relationship = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Edge relationship type (for edge removal)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Cascade = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Cascade delete - remove all connected edges when removing a node",
|
||||
default=True,
|
||||
),
|
||||
]
|
||||
|
||||
ProjectPath = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Project path (defaults to current directory)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class GraphRemoveParams(TypedDict, total=False):
|
||||
"""Parameters for graph remove tool."""
|
||||
|
||||
node_id: Optional[str]
|
||||
source: Optional[str]
|
||||
target: Optional[str]
|
||||
relationship: Optional[str]
|
||||
cascade: bool
|
||||
project_path: Optional[str]
|
||||
|
||||
|
||||
@final
|
||||
class GraphRemoveTool(BaseTool):
|
||||
"""Tool for removing nodes and edges from graph database."""
|
||||
|
||||
def __init__(
|
||||
self, permission_manager: PermissionManager, db_manager: DatabaseManager
|
||||
):
|
||||
"""Initialize the graph remove tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
db_manager: Database manager instance
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.db_manager = db_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "graph_remove"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Remove nodes and edges from the project's graph database.
|
||||
|
||||
To remove a node:
|
||||
- Provide node_id
|
||||
- Use --cascade (default true) to remove connected edges
|
||||
- Use --no-cascade to keep edges (may leave orphaned edges)
|
||||
|
||||
To remove an edge:
|
||||
- Provide source, target, and relationship
|
||||
- Removes only the specific edge
|
||||
|
||||
To remove all edges between two nodes:
|
||||
- Provide source and target (no relationship)
|
||||
|
||||
Examples:
|
||||
- graph_remove --node-id "main.py" # Remove node and its edges
|
||||
- graph_remove --node-id "MyClass" --no-cascade # Remove node only
|
||||
- graph_remove --source "main.py" --target "utils.py" --relationship "imports"
|
||||
- graph_remove --source "func1" --target "func2" # Remove all edges
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[GraphRemoveParams],
|
||||
) -> str:
|
||||
"""Remove nodes or edges from graph.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Result of remove operation
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
node_id = params.get("node_id")
|
||||
source = params.get("source")
|
||||
target = params.get("target")
|
||||
relationship = params.get("relationship")
|
||||
cascade = params.get("cascade", True)
|
||||
project_path = params.get("project_path")
|
||||
|
||||
# Determine if removing node or edge
|
||||
is_node = node_id is not None
|
||||
is_edge = source is not None and target is not None
|
||||
|
||||
if not is_node and not is_edge:
|
||||
return "Error: Must provide either node_id for a node, or (source, target) for edges"
|
||||
|
||||
if is_node and is_edge:
|
||||
return "Error: Cannot remove both node and edge in one operation"
|
||||
|
||||
# Get project database
|
||||
try:
|
||||
if project_path:
|
||||
project_db = self.db_manager.get_project_db(project_path)
|
||||
else:
|
||||
import os
|
||||
|
||||
project_db = self.db_manager.get_project_for_path(os.getcwd())
|
||||
|
||||
if not project_db:
|
||||
return "Error: Could not find project database"
|
||||
|
||||
except PermissionError as e:
|
||||
return str(e)
|
||||
except Exception as e:
|
||||
return f"Error accessing project database: {str(e)}"
|
||||
|
||||
# Get graph connection
|
||||
graph_conn = project_db.get_graph_connection()
|
||||
|
||||
try:
|
||||
if is_node:
|
||||
# Remove node
|
||||
await tool_ctx.info(f"Removing node: {node_id}")
|
||||
|
||||
# Check if node exists
|
||||
cursor = graph_conn.cursor()
|
||||
cursor.execute("SELECT id FROM nodes WHERE id = ?", (node_id,))
|
||||
if not cursor.fetchone():
|
||||
return f"Error: Node '{node_id}' does not exist"
|
||||
|
||||
if cascade:
|
||||
# Count edges that will be removed
|
||||
cursor.execute(
|
||||
"SELECT COUNT(*) FROM edges WHERE source = ? OR target = ?",
|
||||
(node_id, node_id),
|
||||
)
|
||||
edge_count = cursor.fetchone()[0]
|
||||
|
||||
# Remove connected edges
|
||||
graph_conn.execute(
|
||||
"DELETE FROM edges WHERE source = ? OR target = ?",
|
||||
(node_id, node_id),
|
||||
)
|
||||
|
||||
# Remove node
|
||||
graph_conn.execute("DELETE FROM nodes WHERE id = ?", (node_id,))
|
||||
graph_conn.commit()
|
||||
|
||||
# Save to disk
|
||||
project_db._save_graph_to_disk()
|
||||
|
||||
return f"Successfully removed node '{node_id}' and {edge_count} connected edge(s)"
|
||||
else:
|
||||
# Remove node only
|
||||
graph_conn.execute("DELETE FROM nodes WHERE id = ?", (node_id,))
|
||||
graph_conn.commit()
|
||||
|
||||
# Save to disk
|
||||
project_db._save_graph_to_disk()
|
||||
|
||||
return f"Successfully removed node '{node_id}' (edges preserved)"
|
||||
|
||||
else:
|
||||
# Remove edge(s)
|
||||
if relationship:
|
||||
# Remove specific edge
|
||||
await tool_ctx.info(
|
||||
f"Removing edge: {source} --[{relationship}]--> {target}"
|
||||
)
|
||||
|
||||
cursor = graph_conn.cursor()
|
||||
cursor.execute(
|
||||
"DELETE FROM edges WHERE source = ? AND target = ? AND relationship = ?",
|
||||
(source, target, relationship),
|
||||
)
|
||||
|
||||
removed = cursor.rowcount
|
||||
graph_conn.commit()
|
||||
|
||||
if removed == 0:
|
||||
return f"No edge found: {source} --[{relationship}]--> {target}"
|
||||
|
||||
# Save to disk
|
||||
project_db._save_graph_to_disk()
|
||||
|
||||
return f"Successfully removed edge: {source} --[{relationship}]--> {target}"
|
||||
else:
|
||||
# Remove all edges between nodes
|
||||
await tool_ctx.info(
|
||||
f"Removing all edges between {source} and {target}"
|
||||
)
|
||||
|
||||
cursor = graph_conn.cursor()
|
||||
cursor.execute(
|
||||
"DELETE FROM edges WHERE source = ? AND target = ?",
|
||||
(source, target),
|
||||
)
|
||||
|
||||
removed = cursor.rowcount
|
||||
graph_conn.commit()
|
||||
|
||||
if removed == 0:
|
||||
return f"No edges found between '{source}' and '{target}'"
|
||||
|
||||
# Save to disk
|
||||
project_db._save_graph_to_disk()
|
||||
|
||||
return f"Successfully removed {removed} edge(s) between '{source}' and '{target}'"
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to remove from graph: {str(e)}")
|
||||
return f"Error removing from graph: {str(e)}"
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,370 +0,0 @@
|
||||
"""Graph search tool for searching nodes and edges in the graph database."""
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
from typing import Unpack, Optional, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.database.database_manager import DatabaseManager
|
||||
|
||||
Pattern = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Search pattern (SQL LIKE syntax, % for wildcard)",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
SearchType = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="What to search: nodes, edges, properties, all",
|
||||
default="all",
|
||||
),
|
||||
]
|
||||
|
||||
NodeType = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Filter by node type",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Relationship = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Filter by relationship type",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
ProjectPath = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Project path (defaults to current directory)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
MaxResults = Annotated[
|
||||
int,
|
||||
Field(
|
||||
description="Maximum number of results",
|
||||
default=50,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class GraphSearchParams(TypedDict, total=False):
|
||||
"""Parameters for graph search tool."""
|
||||
|
||||
pattern: str
|
||||
search_type: str
|
||||
node_type: Optional[str]
|
||||
relationship: Optional[str]
|
||||
project_path: Optional[str]
|
||||
max_results: int
|
||||
|
||||
|
||||
@final
|
||||
class GraphSearchTool(BaseTool):
|
||||
"""Tool for searching nodes and edges in graph database."""
|
||||
|
||||
def __init__(
|
||||
self, permission_manager: PermissionManager, db_manager: DatabaseManager
|
||||
):
|
||||
"""Initialize the graph search tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
db_manager: Database manager instance
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.db_manager = db_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "graph_search"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Search for nodes and edges in the project's graph database.
|
||||
|
||||
Search types:
|
||||
- nodes: Search in node IDs
|
||||
- edges: Search in edge relationships
|
||||
- properties: Search in node/edge properties
|
||||
- all: Search everywhere (default)
|
||||
|
||||
Supports SQL LIKE pattern matching:
|
||||
- % matches any sequence of characters
|
||||
- _ matches any single character
|
||||
|
||||
Examples:
|
||||
- graph_search --pattern "%test%" # Find anything with 'test'
|
||||
- graph_search --pattern "%.py" --search-type nodes # Find Python files
|
||||
- graph_search --pattern "%import%" --search-type edges
|
||||
- graph_search --pattern "%TODO%" --search-type properties
|
||||
- graph_search --pattern "MyClass%" --node-type "class"
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[GraphSearchParams],
|
||||
) -> str:
|
||||
"""Execute graph search.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Search results
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
pattern = params.get("pattern")
|
||||
if not pattern:
|
||||
return "Error: pattern is required"
|
||||
|
||||
search_type = params.get("search_type", "all")
|
||||
node_type = params.get("node_type")
|
||||
relationship = params.get("relationship")
|
||||
project_path = params.get("project_path")
|
||||
max_results = params.get("max_results", 50)
|
||||
|
||||
# Validate search type
|
||||
valid_types = ["nodes", "edges", "properties", "all"]
|
||||
if search_type not in valid_types:
|
||||
return f"Error: Invalid search_type '{search_type}'. Must be one of: {', '.join(valid_types)}"
|
||||
|
||||
# Get project database
|
||||
try:
|
||||
if project_path:
|
||||
project_db = self.db_manager.get_project_db(project_path)
|
||||
else:
|
||||
import os
|
||||
|
||||
project_db = self.db_manager.get_project_for_path(os.getcwd())
|
||||
|
||||
if not project_db:
|
||||
return "Error: Could not find project database"
|
||||
|
||||
except PermissionError as e:
|
||||
return str(e)
|
||||
except Exception as e:
|
||||
return f"Error accessing project database: {str(e)}"
|
||||
|
||||
await tool_ctx.info(f"Searching graph for pattern: {pattern}")
|
||||
|
||||
# Get graph connection
|
||||
graph_conn = project_db.get_graph_connection()
|
||||
results = []
|
||||
|
||||
try:
|
||||
cursor = graph_conn.cursor()
|
||||
|
||||
# Search nodes
|
||||
if search_type in ["nodes", "all"]:
|
||||
query = "SELECT id, type, properties FROM nodes WHERE id LIKE ?"
|
||||
params_list = [pattern]
|
||||
|
||||
if node_type:
|
||||
query += " AND type = ?"
|
||||
params_list.append(node_type)
|
||||
|
||||
if search_type == "nodes":
|
||||
query += f" LIMIT {max_results}"
|
||||
|
||||
cursor.execute(query, params_list)
|
||||
|
||||
for row in cursor.fetchall():
|
||||
results.append(
|
||||
{
|
||||
"type": "node",
|
||||
"id": row[0],
|
||||
"node_type": row[1],
|
||||
"properties": json.loads(row[2]) if row[2] else {},
|
||||
"match_field": "id",
|
||||
}
|
||||
)
|
||||
|
||||
# Search edges
|
||||
if search_type in ["edges", "all"]:
|
||||
query = """SELECT source, target, relationship, weight, properties
|
||||
FROM edges WHERE relationship LIKE ?"""
|
||||
params_list = [pattern]
|
||||
|
||||
if relationship:
|
||||
query += " AND relationship = ?"
|
||||
params_list.append(relationship)
|
||||
|
||||
if search_type == "edges":
|
||||
query += f" LIMIT {max_results}"
|
||||
|
||||
cursor.execute(query, params_list)
|
||||
|
||||
for row in cursor.fetchall():
|
||||
results.append(
|
||||
{
|
||||
"type": "edge",
|
||||
"source": row[0],
|
||||
"target": row[1],
|
||||
"relationship": row[2],
|
||||
"weight": row[3],
|
||||
"properties": json.loads(row[4]) if row[4] else {},
|
||||
"match_field": "relationship",
|
||||
}
|
||||
)
|
||||
|
||||
# Search in properties
|
||||
if search_type in ["properties", "all"]:
|
||||
# Search node properties
|
||||
query = """SELECT id, type, properties FROM nodes
|
||||
WHERE properties IS NOT NULL AND properties LIKE ?"""
|
||||
params_list = [f"%{pattern}%"]
|
||||
|
||||
if node_type:
|
||||
query += " AND type = ?"
|
||||
params_list.append(node_type)
|
||||
|
||||
cursor.execute(query, params_list)
|
||||
|
||||
for row in cursor.fetchall():
|
||||
props = json.loads(row[2]) if row[2] else {}
|
||||
# Check which property matches
|
||||
matching_props = {}
|
||||
for key, value in props.items():
|
||||
if pattern.replace("%", "").lower() in str(value).lower():
|
||||
matching_props[key] = value
|
||||
|
||||
if matching_props:
|
||||
results.append(
|
||||
{
|
||||
"type": "node",
|
||||
"id": row[0],
|
||||
"node_type": row[1],
|
||||
"properties": props,
|
||||
"match_field": "properties",
|
||||
"matching_properties": matching_props,
|
||||
}
|
||||
)
|
||||
|
||||
# Search edge properties
|
||||
query = """SELECT source, target, relationship, weight, properties
|
||||
FROM edges WHERE properties IS NOT NULL AND properties LIKE ?"""
|
||||
params_list = [f"%{pattern}%"]
|
||||
|
||||
if relationship:
|
||||
query += " AND relationship = ?"
|
||||
params_list.append(relationship)
|
||||
|
||||
cursor.execute(query, params_list)
|
||||
|
||||
for row in cursor.fetchall():
|
||||
props = json.loads(row[4]) if row[4] else {}
|
||||
# Check which property matches
|
||||
matching_props = {}
|
||||
for key, value in props.items():
|
||||
if pattern.replace("%", "").lower() in str(value).lower():
|
||||
matching_props[key] = value
|
||||
|
||||
if matching_props:
|
||||
results.append(
|
||||
{
|
||||
"type": "edge",
|
||||
"source": row[0],
|
||||
"target": row[1],
|
||||
"relationship": row[2],
|
||||
"weight": row[3],
|
||||
"properties": props,
|
||||
"match_field": "properties",
|
||||
"matching_properties": matching_props,
|
||||
}
|
||||
)
|
||||
|
||||
# Limit total results if searching all
|
||||
if search_type == "all" and len(results) > max_results:
|
||||
results = results[:max_results]
|
||||
|
||||
if not results:
|
||||
return f"No results found for pattern '{pattern}'"
|
||||
|
||||
# Format results
|
||||
output = [f"Found {len(results)} result(s) for pattern '{pattern}':\n"]
|
||||
|
||||
# Group by type
|
||||
nodes = [r for r in results if r["type"] == "node"]
|
||||
edges = [r for r in results if r["type"] == "edge"]
|
||||
|
||||
if nodes:
|
||||
output.append(f"Nodes ({len(nodes)}):")
|
||||
for node in nodes[:20]: # Show first 20
|
||||
output.append(f" {node['id']} ({node['node_type']})")
|
||||
if (
|
||||
node["match_field"] == "properties"
|
||||
and "matching_properties" in node
|
||||
):
|
||||
output.append(
|
||||
f" Matched in: {list(node['matching_properties'].keys())}"
|
||||
)
|
||||
if node["properties"] and node["match_field"] != "properties":
|
||||
props_str = json.dumps(node["properties"], indent=6)[:100]
|
||||
if len(props_str) == 100:
|
||||
props_str += "..."
|
||||
output.append(f" Properties: {props_str}")
|
||||
|
||||
if len(nodes) > 20:
|
||||
output.append(f" ... and {len(nodes) - 20} more nodes")
|
||||
output.append("")
|
||||
|
||||
if edges:
|
||||
output.append(f"Edges ({len(edges)}):")
|
||||
for edge in edges[:20]: # Show first 20
|
||||
output.append(
|
||||
f" {edge['source']} --[{edge['relationship']}]--> {edge['target']}"
|
||||
)
|
||||
if (
|
||||
edge["match_field"] == "properties"
|
||||
and "matching_properties" in edge
|
||||
):
|
||||
output.append(
|
||||
f" Matched in: {list(edge['matching_properties'].keys())}"
|
||||
)
|
||||
if edge["weight"] != 1.0:
|
||||
output.append(f" Weight: {edge['weight']}")
|
||||
if edge["properties"]:
|
||||
props_str = json.dumps(edge["properties"], indent=6)[:100]
|
||||
if len(props_str) == 100:
|
||||
props_str += "..."
|
||||
output.append(f" Properties: {props_str}")
|
||||
|
||||
if len(edges) > 20:
|
||||
output.append(f" ... and {len(edges) - 20} more edges")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except sqlite3.Error as e:
|
||||
await tool_ctx.error(f"SQL error: {str(e)}")
|
||||
return f"SQL error: {str(e)}"
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Unexpected error: {str(e)}")
|
||||
return f"Error executing search: {str(e)}"
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,377 +0,0 @@
|
||||
"""Graph statistics tool for analyzing the graph database."""
|
||||
|
||||
import sqlite3
|
||||
from typing import Unpack, Optional, Annotated, TypedDict, final, override
|
||||
from collections import defaultdict
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.database.database_manager import DatabaseManager
|
||||
|
||||
ProjectPath = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Project path (defaults to current directory)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Detailed = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Show detailed statistics",
|
||||
default=False,
|
||||
),
|
||||
]
|
||||
|
||||
NodeType = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Filter stats by node type",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Relationship = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Filter stats by relationship type",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class GraphStatsParams(TypedDict, total=False):
|
||||
"""Parameters for graph stats tool."""
|
||||
|
||||
project_path: Optional[str]
|
||||
detailed: bool
|
||||
node_type: Optional[str]
|
||||
relationship: Optional[str]
|
||||
|
||||
|
||||
@final
|
||||
class GraphStatsTool(BaseTool):
|
||||
"""Tool for getting graph database statistics."""
|
||||
|
||||
def __init__(
|
||||
self, permission_manager: PermissionManager, db_manager: DatabaseManager
|
||||
):
|
||||
"""Initialize the graph stats tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
db_manager: Database manager instance
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.db_manager = db_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "graph_stats"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Get statistics about the project's graph database.
|
||||
|
||||
Shows:
|
||||
- Node and edge counts
|
||||
- Node type distribution
|
||||
- Relationship type distribution
|
||||
- Degree statistics (connections per node)
|
||||
- Connected components
|
||||
- Most connected nodes (hubs)
|
||||
- Orphaned nodes
|
||||
|
||||
Examples:
|
||||
- graph_stats # Basic stats
|
||||
- graph_stats --detailed # Detailed analysis
|
||||
- graph_stats --node-type "class" # Stats for specific node type
|
||||
- graph_stats --relationship "calls" # Stats for specific relationship
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[GraphStatsParams],
|
||||
) -> str:
|
||||
"""Get graph statistics.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Graph statistics
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
project_path = params.get("project_path")
|
||||
detailed = params.get("detailed", False)
|
||||
node_type_filter = params.get("node_type")
|
||||
relationship_filter = params.get("relationship")
|
||||
|
||||
# Get project database
|
||||
try:
|
||||
if project_path:
|
||||
project_db = self.db_manager.get_project_db(project_path)
|
||||
else:
|
||||
import os
|
||||
|
||||
project_db = self.db_manager.get_project_for_path(os.getcwd())
|
||||
|
||||
if not project_db:
|
||||
return "Error: Could not find project database"
|
||||
|
||||
except PermissionError as e:
|
||||
return str(e)
|
||||
except Exception as e:
|
||||
return f"Error accessing project database: {str(e)}"
|
||||
|
||||
await tool_ctx.info(
|
||||
f"Getting graph statistics for project: {project_db.project_path}"
|
||||
)
|
||||
|
||||
# Get graph connection
|
||||
graph_conn = project_db.get_graph_connection()
|
||||
|
||||
try:
|
||||
cursor = graph_conn.cursor()
|
||||
output = []
|
||||
output.append(f"=== Graph Database Statistics ===")
|
||||
output.append(f"Project: {project_db.project_path}")
|
||||
output.append(f"Database: {project_db.graph_path}")
|
||||
output.append("")
|
||||
|
||||
# Basic counts
|
||||
if node_type_filter:
|
||||
cursor.execute(
|
||||
"SELECT COUNT(*) FROM nodes WHERE type = ?", (node_type_filter,)
|
||||
)
|
||||
node_count = cursor.fetchone()[0]
|
||||
output.append(f"Nodes (type='{node_type_filter}'): {node_count:,}")
|
||||
else:
|
||||
cursor.execute("SELECT COUNT(*) FROM nodes")
|
||||
node_count = cursor.fetchone()[0]
|
||||
output.append(f"Total Nodes: {node_count:,}")
|
||||
|
||||
if relationship_filter:
|
||||
cursor.execute(
|
||||
"SELECT COUNT(*) FROM edges WHERE relationship = ?",
|
||||
(relationship_filter,),
|
||||
)
|
||||
edge_count = cursor.fetchone()[0]
|
||||
output.append(
|
||||
f"Edges (relationship='{relationship_filter}'): {edge_count:,}"
|
||||
)
|
||||
else:
|
||||
cursor.execute("SELECT COUNT(*) FROM edges")
|
||||
edge_count = cursor.fetchone()[0]
|
||||
output.append(f"Total Edges: {edge_count:,}")
|
||||
|
||||
if node_count == 0:
|
||||
output.append("\nGraph is empty.")
|
||||
return "\n".join(output)
|
||||
|
||||
output.append("")
|
||||
|
||||
# Node type distribution
|
||||
output.append("=== Node Types ===")
|
||||
cursor.execute(
|
||||
"SELECT type, COUNT(*) as count FROM nodes GROUP BY type ORDER BY count DESC"
|
||||
)
|
||||
node_types = cursor.fetchall()
|
||||
|
||||
for n_type, count in node_types[:10]:
|
||||
pct = (count / node_count) * 100
|
||||
output.append(f"{n_type}: {count:,} ({pct:.1f}%)")
|
||||
|
||||
if len(node_types) > 10:
|
||||
output.append(f"... and {len(node_types) - 10} more types")
|
||||
|
||||
output.append("")
|
||||
|
||||
# Relationship distribution
|
||||
output.append("=== Relationship Types ===")
|
||||
cursor.execute(
|
||||
"SELECT relationship, COUNT(*) as count FROM edges GROUP BY relationship ORDER BY count DESC"
|
||||
)
|
||||
rel_types = cursor.fetchall()
|
||||
|
||||
if rel_types:
|
||||
for rel, count in rel_types[:10]:
|
||||
pct = (count / edge_count) * 100 if edge_count > 0 else 0
|
||||
output.append(f"{rel}: {count:,} ({pct:.1f}%)")
|
||||
|
||||
if len(rel_types) > 10:
|
||||
output.append(f"... and {len(rel_types) - 10} more types")
|
||||
else:
|
||||
output.append("No edges in graph")
|
||||
|
||||
output.append("")
|
||||
|
||||
# Degree statistics
|
||||
output.append("=== Connectivity ===")
|
||||
|
||||
# Calculate degrees
|
||||
degrees = defaultdict(int)
|
||||
|
||||
# Out-degree
|
||||
query = "SELECT source, COUNT(*) FROM edges"
|
||||
if relationship_filter:
|
||||
query += " WHERE relationship = ?"
|
||||
cursor.execute(query + " GROUP BY source", (relationship_filter,))
|
||||
else:
|
||||
cursor.execute(query + " GROUP BY source")
|
||||
|
||||
for node, out_degree in cursor.fetchall():
|
||||
degrees[node] += out_degree
|
||||
|
||||
# In-degree
|
||||
query = "SELECT target, COUNT(*) FROM edges"
|
||||
if relationship_filter:
|
||||
query += " WHERE relationship = ?"
|
||||
cursor.execute(query + " GROUP BY target", (relationship_filter,))
|
||||
else:
|
||||
cursor.execute(query + " GROUP BY target")
|
||||
|
||||
for node, in_degree in cursor.fetchall():
|
||||
degrees[node] += in_degree
|
||||
|
||||
if degrees:
|
||||
degree_values = list(degrees.values())
|
||||
avg_degree = sum(degree_values) / len(degree_values)
|
||||
max_degree = max(degree_values)
|
||||
min_degree = min(degree_values)
|
||||
|
||||
output.append(f"Average degree: {avg_degree:.2f}")
|
||||
output.append(f"Max degree: {max_degree}")
|
||||
output.append(f"Min degree: {min_degree}")
|
||||
|
||||
# Most connected nodes
|
||||
output.append("\nMost connected nodes:")
|
||||
sorted_nodes = sorted(degrees.items(), key=lambda x: x[1], reverse=True)
|
||||
|
||||
for node, degree in sorted_nodes[:5]:
|
||||
cursor.execute("SELECT type FROM nodes WHERE id = ?", (node,))
|
||||
node_type = cursor.fetchone()
|
||||
type_str = f" ({node_type[0]})" if node_type else ""
|
||||
output.append(f" {node}{type_str}: {degree} connections")
|
||||
|
||||
# Orphaned nodes
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT COUNT(*) FROM nodes n
|
||||
WHERE NOT EXISTS (SELECT 1 FROM edges WHERE source = n.id OR target = n.id)
|
||||
"""
|
||||
)
|
||||
orphan_count = cursor.fetchone()[0]
|
||||
if orphan_count > 0:
|
||||
orphan_pct = (orphan_count / node_count) * 100
|
||||
output.append(f"\nOrphaned nodes: {orphan_count} ({orphan_pct:.1f}%)")
|
||||
|
||||
if detailed:
|
||||
output.append("\n=== Detailed Analysis ===")
|
||||
|
||||
# Node properties usage
|
||||
cursor.execute(
|
||||
"SELECT COUNT(*) FROM nodes WHERE properties IS NOT NULL"
|
||||
)
|
||||
nodes_with_props = cursor.fetchone()[0]
|
||||
if nodes_with_props > 0:
|
||||
props_pct = (nodes_with_props / node_count) * 100
|
||||
output.append(
|
||||
f"Nodes with properties: {nodes_with_props} ({props_pct:.1f}%)"
|
||||
)
|
||||
|
||||
# Edge properties usage
|
||||
cursor.execute(
|
||||
"SELECT COUNT(*) FROM edges WHERE properties IS NOT NULL"
|
||||
)
|
||||
edges_with_props = cursor.fetchone()[0]
|
||||
if edges_with_props > 0 and edge_count > 0:
|
||||
props_pct = (edges_with_props / edge_count) * 100
|
||||
output.append(
|
||||
f"Edges with properties: {edges_with_props} ({props_pct:.1f}%)"
|
||||
)
|
||||
|
||||
# Weight distribution
|
||||
cursor.execute(
|
||||
"SELECT MIN(weight), MAX(weight), AVG(weight) FROM edges"
|
||||
)
|
||||
weight_stats = cursor.fetchone()
|
||||
if weight_stats[0] is not None:
|
||||
output.append(f"\nEdge weights:")
|
||||
output.append(f" Min: {weight_stats[0]}")
|
||||
output.append(f" Max: {weight_stats[1]}")
|
||||
output.append(f" Avg: {weight_stats[2]:.2f}")
|
||||
|
||||
# Most common patterns
|
||||
if not relationship_filter:
|
||||
output.append("\n=== Common Patterns ===")
|
||||
|
||||
# Most common node type connections
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT n1.type, e.relationship, n2.type, COUNT(*) as count
|
||||
FROM edges e
|
||||
JOIN nodes n1 ON e.source = n1.id
|
||||
JOIN nodes n2 ON e.target = n2.id
|
||||
GROUP BY n1.type, e.relationship, n2.type
|
||||
ORDER BY count DESC
|
||||
LIMIT 10
|
||||
"""
|
||||
)
|
||||
|
||||
patterns = cursor.fetchall()
|
||||
if patterns:
|
||||
output.append("Most common connections:")
|
||||
for src_type, rel, tgt_type, count in patterns:
|
||||
output.append(
|
||||
f" {src_type} --[{rel}]--> {tgt_type}: {count} times"
|
||||
)
|
||||
|
||||
# Component analysis (simplified)
|
||||
output.append("\n=== Graph Structure ===")
|
||||
|
||||
# Check if graph is fully connected (simplified)
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT COUNT(DISTINCT node_id) FROM (
|
||||
SELECT source as node_id FROM edges
|
||||
UNION
|
||||
SELECT target as node_id FROM edges
|
||||
)
|
||||
"""
|
||||
)
|
||||
connected_nodes = cursor.fetchone()[0]
|
||||
|
||||
if connected_nodes < node_count:
|
||||
output.append(f"Connected nodes: {connected_nodes} / {node_count}")
|
||||
output.append("Graph has disconnected components")
|
||||
else:
|
||||
output.append("All nodes are connected")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except sqlite3.Error as e:
|
||||
await tool_ctx.error(f"SQL error: {str(e)}")
|
||||
return f"SQL error: {str(e)}"
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Unexpected error: {str(e)}")
|
||||
return f"Error getting statistics: {str(e)}"
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,439 +0,0 @@
|
||||
"""Unified SQL database tool."""
|
||||
|
||||
import sqlite3
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
Unpack,
|
||||
Optional,
|
||||
Annotated,
|
||||
TypedDict,
|
||||
final,
|
||||
override,
|
||||
)
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.database.database_manager import DatabaseManager
|
||||
|
||||
# Parameter types
|
||||
Query = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="SQL query to execute",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Pattern = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Search pattern for table/column names or data",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Table = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Table name for operations",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Action = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Action: query (default), search, schema, stats",
|
||||
default="query",
|
||||
),
|
||||
]
|
||||
|
||||
Limit = Annotated[
|
||||
int,
|
||||
Field(
|
||||
description="Maximum rows to return",
|
||||
default=100,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class SQLParams(TypedDict, total=False):
|
||||
"""Parameters for SQL tool."""
|
||||
|
||||
query: Optional[str]
|
||||
pattern: Optional[str]
|
||||
table: Optional[str]
|
||||
action: str
|
||||
limit: int
|
||||
|
||||
|
||||
@final
|
||||
class SQLTool(BaseTool):
|
||||
"""Unified SQL database tool."""
|
||||
|
||||
def __init__(
|
||||
self, permission_manager: PermissionManager, db_manager: DatabaseManager
|
||||
):
|
||||
"""Initialize the SQL tool."""
|
||||
super().__init__(permission_manager)
|
||||
self.db_manager = db_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "sql"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """SQLite database. Actions: query (default), search, schema, stats.
|
||||
|
||||
Usage:
|
||||
sql "SELECT * FROM users WHERE active = 1"
|
||||
sql --action schema
|
||||
sql --action search --pattern "john"
|
||||
sql --action stats --table users
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[SQLParams],
|
||||
) -> str:
|
||||
"""Execute SQL operation."""
|
||||
tool_ctx = self.create_tool_context(ctx)
|
||||
|
||||
# Get current project database
|
||||
project_db = self.db_manager.get_current_project_db()
|
||||
if not project_db:
|
||||
return "Error: No project database found. Are you in a project directory?"
|
||||
|
||||
# Extract action
|
||||
action = params.get("action", "query")
|
||||
|
||||
# Route to appropriate handler
|
||||
if action == "query":
|
||||
return await self._handle_query(project_db, params, tool_ctx)
|
||||
elif action == "search":
|
||||
return await self._handle_search(project_db, params, tool_ctx)
|
||||
elif action == "schema":
|
||||
return await self._handle_schema(project_db, params, tool_ctx)
|
||||
elif action == "stats":
|
||||
return await self._handle_stats(project_db, params, tool_ctx)
|
||||
else:
|
||||
return f"Error: Unknown action '{action}'. Valid actions: query, search, schema, stats"
|
||||
|
||||
async def _handle_query(self, project_db, params: Dict[str, Any], tool_ctx) -> str:
|
||||
"""Execute SQL query."""
|
||||
query = params.get("query")
|
||||
if not query:
|
||||
return "Error: query required for query action"
|
||||
|
||||
limit = params.get("limit", 100)
|
||||
|
||||
try:
|
||||
with project_db.get_sqlite_connection() as conn:
|
||||
# Enable row factory for dict-like access
|
||||
conn.row_factory = sqlite3.Row
|
||||
|
||||
# Add LIMIT if not present in SELECT queries
|
||||
query_upper = query.upper().strip()
|
||||
if query_upper.startswith("SELECT") and "LIMIT" not in query_upper:
|
||||
query = f"{query} LIMIT {limit}"
|
||||
|
||||
cursor = conn.execute(query)
|
||||
|
||||
# Handle different query types
|
||||
if query_upper.startswith("SELECT"):
|
||||
rows = cursor.fetchall()
|
||||
|
||||
if not rows:
|
||||
return "No results found"
|
||||
|
||||
# Get column names
|
||||
columns = [description[0] for description in cursor.description]
|
||||
|
||||
# Format as table
|
||||
output = ["=== Query Results ==="]
|
||||
output.append(f"Columns: {', '.join(columns)}")
|
||||
output.append("-" * 60)
|
||||
|
||||
for row in rows:
|
||||
row_data = []
|
||||
for col in columns:
|
||||
value = row[col]
|
||||
if value is None:
|
||||
value = "NULL"
|
||||
elif isinstance(value, str) and len(value) > 50:
|
||||
value = value[:50] + "..."
|
||||
row_data.append(str(value))
|
||||
output.append(" | ".join(row_data))
|
||||
|
||||
output.append(f"\nRows returned: {len(rows)}")
|
||||
if len(rows) == limit:
|
||||
output.append(f"(Limited to {limit} rows)")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
else:
|
||||
# For INSERT, UPDATE, DELETE
|
||||
conn.commit()
|
||||
rows_affected = cursor.rowcount
|
||||
|
||||
if query_upper.startswith("INSERT"):
|
||||
return f"Inserted {rows_affected} row(s)"
|
||||
elif query_upper.startswith("UPDATE"):
|
||||
return f"Updated {rows_affected} row(s)"
|
||||
elif query_upper.startswith("DELETE"):
|
||||
return f"Deleted {rows_affected} row(s)"
|
||||
else:
|
||||
return f"Query executed successfully. Rows affected: {rows_affected}"
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Query failed: {str(e)}")
|
||||
return f"Error executing query: {str(e)}"
|
||||
|
||||
async def _handle_search(self, project_db, params: Dict[str, Any], tool_ctx) -> str:
|
||||
"""Search for data in tables."""
|
||||
pattern = params.get("pattern")
|
||||
if not pattern:
|
||||
return "Error: pattern required for search action"
|
||||
|
||||
table = params.get("table")
|
||||
limit = params.get("limit", 100)
|
||||
|
||||
try:
|
||||
with project_db.get_sqlite_connection() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
|
||||
# Get all tables if not specified
|
||||
if not table:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
SELECT name FROM sqlite_master
|
||||
WHERE type='table' AND name NOT LIKE 'sqlite_%'
|
||||
"""
|
||||
)
|
||||
tables = [row[0] for row in cursor.fetchall()]
|
||||
else:
|
||||
tables = [table]
|
||||
|
||||
all_results = []
|
||||
|
||||
for tbl in tables:
|
||||
# Get columns
|
||||
cursor = conn.execute(f"PRAGMA table_info({tbl})")
|
||||
columns = [row[1] for row in cursor.fetchall()]
|
||||
|
||||
# Build search query
|
||||
where_clauses = [f"{col} LIKE ?" for col in columns]
|
||||
query = f"SELECT * FROM {tbl} WHERE {' OR '.join(where_clauses)} LIMIT {limit}"
|
||||
|
||||
# Search
|
||||
cursor = conn.execute(query, [f"%{pattern}%"] * len(columns))
|
||||
rows = cursor.fetchall()
|
||||
|
||||
if rows:
|
||||
all_results.append((tbl, columns, rows))
|
||||
|
||||
if not all_results:
|
||||
return f"No results found for pattern '{pattern}'"
|
||||
|
||||
# Format results
|
||||
output = [f"=== Search Results for '{pattern}' ==="]
|
||||
|
||||
for tbl, columns, rows in all_results:
|
||||
output.append(f"\nTable: {tbl}")
|
||||
output.append(f"Columns: {', '.join(columns)}")
|
||||
output.append("-" * 60)
|
||||
|
||||
for row in rows:
|
||||
row_data = []
|
||||
for col in columns:
|
||||
value = row[col]
|
||||
if value is None:
|
||||
value = "NULL"
|
||||
elif isinstance(value, str):
|
||||
# Highlight matches
|
||||
if pattern.lower() in str(value).lower():
|
||||
value = f"**{value}**"
|
||||
if len(value) > 50:
|
||||
value = value[:50] + "..."
|
||||
row_data.append(str(value))
|
||||
output.append(" | ".join(row_data))
|
||||
|
||||
output.append(f"Found {len(rows)} row(s) in {tbl}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Search failed: {str(e)}")
|
||||
return f"Error during search: {str(e)}"
|
||||
|
||||
async def _handle_schema(self, project_db, params: Dict[str, Any], tool_ctx) -> str:
|
||||
"""Show database schema."""
|
||||
table = params.get("table")
|
||||
|
||||
try:
|
||||
with project_db.get_sqlite_connection() as conn:
|
||||
if table:
|
||||
# Show specific table schema
|
||||
cursor = conn.execute(f"PRAGMA table_info({table})")
|
||||
columns = cursor.fetchall()
|
||||
|
||||
if not columns:
|
||||
return f"Table '{table}' not found"
|
||||
|
||||
output = [f"=== Schema for table '{table}' ==="]
|
||||
output.append("Column | Type | Not Null | Default | Primary Key")
|
||||
output.append("-" * 60)
|
||||
|
||||
for col in columns:
|
||||
output.append(
|
||||
f"{col[1]} | {col[2]} | {col[3]} | {col[4]} | {col[5]}"
|
||||
)
|
||||
|
||||
# Get indexes
|
||||
cursor = conn.execute(f"PRAGMA index_list({table})")
|
||||
indexes = cursor.fetchall()
|
||||
|
||||
if indexes:
|
||||
output.append("\nIndexes:")
|
||||
for idx in indexes:
|
||||
output.append(f" {idx[1]} (unique: {idx[2]})")
|
||||
|
||||
else:
|
||||
# Show all tables
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
SELECT name, sql FROM sqlite_master
|
||||
WHERE type='table' AND name NOT LIKE 'sqlite_%'
|
||||
ORDER BY name
|
||||
"""
|
||||
)
|
||||
tables = cursor.fetchall()
|
||||
|
||||
if not tables:
|
||||
return "No tables found in database"
|
||||
|
||||
output = ["=== Database Schema ==="]
|
||||
|
||||
for table_name, _create_sql in tables:
|
||||
output.append(f"\nTable: {table_name}")
|
||||
|
||||
# Get row count
|
||||
cursor = conn.execute(f"SELECT COUNT(*) FROM {table_name}")
|
||||
count = cursor.fetchone()[0]
|
||||
output.append(f"Rows: {count}")
|
||||
|
||||
# Get columns
|
||||
cursor = conn.execute(f"PRAGMA table_info({table_name})")
|
||||
columns = cursor.fetchall()
|
||||
output.append(
|
||||
f"Columns: {', '.join([col[1] for col in columns])}"
|
||||
)
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to get schema: {str(e)}")
|
||||
return f"Error getting schema: {str(e)}"
|
||||
|
||||
async def _handle_stats(self, project_db, params: Dict[str, Any], tool_ctx) -> str:
|
||||
"""Get database statistics."""
|
||||
table = params.get("table")
|
||||
|
||||
try:
|
||||
with project_db.get_sqlite_connection() as conn:
|
||||
output = ["=== Database Statistics ==="]
|
||||
output.append(f"Database: {project_db.sqlite_path}")
|
||||
|
||||
# Get file size
|
||||
db_size = project_db.sqlite_path.stat().st_size
|
||||
output.append(f"Size: {db_size / 1024 / 1024:.2f} MB")
|
||||
|
||||
if table:
|
||||
# Stats for specific table
|
||||
cursor = conn.execute(f"SELECT COUNT(*) FROM {table}")
|
||||
count = cursor.fetchone()[0]
|
||||
|
||||
output.append(f"\nTable: {table}")
|
||||
output.append(f"Total rows: {count}")
|
||||
|
||||
# Get column stats
|
||||
cursor = conn.execute(f"PRAGMA table_info({table})")
|
||||
columns = cursor.fetchall()
|
||||
|
||||
output.append("\nColumn statistics:")
|
||||
for col in columns:
|
||||
col_name = col[1]
|
||||
col_type = col[2]
|
||||
|
||||
# Get basic stats based on type
|
||||
if "INT" in col_type.upper() or "REAL" in col_type.upper():
|
||||
cursor = conn.execute(
|
||||
f"""
|
||||
SELECT
|
||||
MIN({col_name}) as min_val,
|
||||
MAX({col_name}) as max_val,
|
||||
AVG({col_name}) as avg_val,
|
||||
COUNT(DISTINCT {col_name}) as distinct_count
|
||||
FROM {table}
|
||||
"""
|
||||
)
|
||||
stats = cursor.fetchone()
|
||||
output.append(
|
||||
f" {col_name}: min={stats[0]}, max={stats[1]}, avg={stats[2]:.2f}, distinct={stats[3]}"
|
||||
)
|
||||
else:
|
||||
cursor = conn.execute(
|
||||
f"""
|
||||
SELECT
|
||||
COUNT(DISTINCT {col_name}) as distinct_count,
|
||||
COUNT(*) - COUNT({col_name}) as null_count
|
||||
FROM {table}
|
||||
"""
|
||||
)
|
||||
stats = cursor.fetchone()
|
||||
output.append(
|
||||
f" {col_name}: distinct={stats[0]}, nulls={stats[1]}"
|
||||
)
|
||||
|
||||
else:
|
||||
# Overall database stats
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
SELECT name FROM sqlite_master
|
||||
WHERE type='table' AND name NOT LIKE 'sqlite_%'
|
||||
"""
|
||||
)
|
||||
tables = cursor.fetchall()
|
||||
|
||||
output.append(f"\nTotal tables: {len(tables)}")
|
||||
output.append("\nTable row counts:")
|
||||
|
||||
total_rows = 0
|
||||
for (table_name,) in tables:
|
||||
cursor = conn.execute(f"SELECT COUNT(*) FROM {table_name}")
|
||||
count = cursor.fetchone()[0]
|
||||
total_rows += count
|
||||
output.append(f" {table_name}: {count} rows")
|
||||
|
||||
output.append(f"\nTotal rows across all tables: {total_rows}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to get stats: {str(e)}")
|
||||
return f"Error getting stats: {str(e)}"
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,235 +0,0 @@
|
||||
"""SQL query tool for direct database queries."""
|
||||
|
||||
import sqlite3
|
||||
from typing import Unpack, Optional, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.database.database_manager import DatabaseManager
|
||||
|
||||
Query = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="SQL query to execute",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
ProjectPath = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Project path (defaults to current directory)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
ReadOnly = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Execute in read-only mode (no INSERT/UPDATE/DELETE)",
|
||||
default=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class SqlQueryParams(TypedDict, total=False):
|
||||
"""Parameters for SQL query tool."""
|
||||
|
||||
query: str
|
||||
project_path: Optional[str]
|
||||
read_only: bool
|
||||
|
||||
|
||||
@final
|
||||
class SqlQueryTool(BaseTool):
|
||||
"""Tool for executing SQL queries on project databases."""
|
||||
|
||||
def __init__(
|
||||
self, permission_manager: PermissionManager, db_manager: DatabaseManager
|
||||
):
|
||||
"""Initialize the SQL query tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
db_manager: Database manager instance
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.db_manager = db_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "sql_query"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Execute SQL queries on the project's embedded SQLite database.
|
||||
|
||||
Each project has its own SQLite database with tables:
|
||||
- metadata: Key-value store for project metadata
|
||||
- files: File information and content
|
||||
- symbols: Code symbols (functions, classes, etc.)
|
||||
|
||||
Features:
|
||||
- Direct SQL query execution
|
||||
- Read-only mode by default (safety)
|
||||
- Returns results in tabular format
|
||||
- Automatic project detection
|
||||
|
||||
Examples:
|
||||
- sql_query --query "SELECT * FROM files LIMIT 10"
|
||||
- sql_query --query "SELECT name, type FROM symbols WHERE type='function'"
|
||||
- sql_query --query "INSERT INTO metadata (key, value) VALUES ('version', '1.0')" --read-only false
|
||||
|
||||
Note: Use sql_search for text search operations."""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[SqlQueryParams],
|
||||
) -> str:
|
||||
"""Execute SQL query.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Query results
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
query = params.get("query")
|
||||
if not query:
|
||||
return "Error: query is required"
|
||||
|
||||
project_path = params.get("project_path")
|
||||
read_only = params.get("read_only", True)
|
||||
|
||||
# Get project database
|
||||
try:
|
||||
if project_path:
|
||||
project_db = self.db_manager.get_project_db(project_path)
|
||||
else:
|
||||
import os
|
||||
|
||||
project_db = self.db_manager.get_project_for_path(os.getcwd())
|
||||
|
||||
if not project_db:
|
||||
return "Error: Could not find project database"
|
||||
|
||||
except PermissionError as e:
|
||||
return str(e)
|
||||
except Exception as e:
|
||||
return f"Error accessing project database: {str(e)}"
|
||||
|
||||
# Check if query is read-only
|
||||
if read_only:
|
||||
# Simple check for write operations
|
||||
write_keywords = ["INSERT", "UPDATE", "DELETE", "DROP", "CREATE", "ALTER"]
|
||||
query_upper = query.upper()
|
||||
for keyword in write_keywords:
|
||||
if keyword in query_upper:
|
||||
return f"Error: Query contains {keyword} operation. Set --read-only false to allow write operations."
|
||||
|
||||
await tool_ctx.info(
|
||||
f"Executing SQL query on project: {project_db.project_path}"
|
||||
)
|
||||
|
||||
# Execute query
|
||||
conn = None
|
||||
try:
|
||||
conn = project_db.get_sqlite_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Execute the query
|
||||
cursor.execute(query)
|
||||
|
||||
# Handle different query types
|
||||
if query.strip().upper().startswith("SELECT"):
|
||||
# Fetch results
|
||||
results = cursor.fetchall()
|
||||
|
||||
if not results:
|
||||
return "No results found."
|
||||
|
||||
# Get column names
|
||||
columns = [desc[0] for desc in cursor.description]
|
||||
|
||||
# Format as table
|
||||
output = self._format_results_table(columns, results)
|
||||
|
||||
return f"Query executed successfully. Found {len(results)} row(s).\n\n{output}"
|
||||
|
||||
else:
|
||||
# For non-SELECT queries, commit and return affected rows
|
||||
conn.commit()
|
||||
affected = cursor.rowcount
|
||||
return f"Query executed successfully. Affected {affected} row(s)."
|
||||
|
||||
except sqlite3.Error as e:
|
||||
await tool_ctx.error(f"SQL error: {str(e)}")
|
||||
return f"SQL error: {str(e)}"
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Unexpected error: {str(e)}")
|
||||
return f"Error executing query: {str(e)}"
|
||||
finally:
|
||||
if conn:
|
||||
conn.close()
|
||||
|
||||
def _format_results_table(self, columns: list[str], rows: list[tuple]) -> str:
|
||||
"""Format query results as a table."""
|
||||
if not rows:
|
||||
return "No results"
|
||||
|
||||
# Calculate column widths
|
||||
col_widths = []
|
||||
for i, col in enumerate(columns):
|
||||
max_width = len(col)
|
||||
for row in rows[:100]: # Check first 100 rows
|
||||
val_str = str(row[i]) if row[i] is not None else "NULL"
|
||||
max_width = max(max_width, len(val_str))
|
||||
col_widths.append(min(max_width, 50)) # Cap at 50 chars
|
||||
|
||||
# Build header
|
||||
header = " | ".join(col.ljust(width) for col, width in zip(columns, col_widths))
|
||||
separator = "-+-".join("-" * width for width in col_widths)
|
||||
|
||||
# Build rows
|
||||
output_rows = []
|
||||
for row in rows[:1000]: # Limit to 1000 rows
|
||||
row_str = " | ".join(
|
||||
self._truncate(str(val) if val is not None else "NULL", width).ljust(
|
||||
width
|
||||
)
|
||||
for val, width in zip(row, col_widths)
|
||||
)
|
||||
output_rows.append(row_str)
|
||||
|
||||
# Combine
|
||||
output = [header, separator] + output_rows
|
||||
|
||||
if len(rows) > 1000:
|
||||
output.append(f"\n... and {len(rows) - 1000} more rows")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _truncate(self, text: str, max_width: int) -> str:
|
||||
"""Truncate text to max width."""
|
||||
if len(text) <= max_width:
|
||||
return text
|
||||
return text[: max_width - 3] + "..."
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,300 +0,0 @@
|
||||
"""SQL search tool for text search in database."""
|
||||
|
||||
import sqlite3
|
||||
from typing import Unpack, Optional, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.database.database_manager import DatabaseManager
|
||||
|
||||
SearchPattern = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Search pattern (SQL LIKE syntax, % for wildcard)",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
Table = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Table to search in (files, symbols, metadata)",
|
||||
default="files",
|
||||
),
|
||||
]
|
||||
|
||||
Column = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Specific column to search (searches all text columns if not specified)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
ProjectPath = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Project path (defaults to current directory)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
MaxResults = Annotated[
|
||||
int,
|
||||
Field(
|
||||
description="Maximum number of results",
|
||||
default=50,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class SqlSearchParams(TypedDict, total=False):
|
||||
"""Parameters for SQL search tool."""
|
||||
|
||||
pattern: str
|
||||
table: str
|
||||
column: Optional[str]
|
||||
project_path: Optional[str]
|
||||
max_results: int
|
||||
|
||||
|
||||
@final
|
||||
class SqlSearchTool(BaseTool):
|
||||
"""Tool for searching text in SQLite database."""
|
||||
|
||||
def __init__(
|
||||
self, permission_manager: PermissionManager, db_manager: DatabaseManager
|
||||
):
|
||||
"""Initialize the SQL search tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
db_manager: Database manager instance
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.db_manager = db_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "sql_search"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Search for text patterns in the project's SQLite database.
|
||||
|
||||
Supports SQL LIKE pattern matching:
|
||||
- % matches any sequence of characters
|
||||
- _ matches any single character
|
||||
- Use %pattern% to search for pattern anywhere
|
||||
|
||||
Tables available:
|
||||
- files: Search in file paths and content
|
||||
- symbols: Search in symbol names, types, and signatures
|
||||
- metadata: Search in key-value metadata
|
||||
|
||||
Examples:
|
||||
- sql_search --pattern "%TODO%" --table files
|
||||
- sql_search --pattern "test_%" --table symbols --column name
|
||||
- sql_search --pattern "%config%" --table metadata
|
||||
- sql_search --pattern "%.py" --table files --column path
|
||||
|
||||
Use sql_query for complex queries with joins, conditions, etc."""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[SqlSearchParams],
|
||||
) -> str:
|
||||
"""Execute SQL search.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Search results
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
pattern = params.get("pattern")
|
||||
if not pattern:
|
||||
return "Error: pattern is required"
|
||||
|
||||
table = params.get("table", "files")
|
||||
column = params.get("column")
|
||||
project_path = params.get("project_path")
|
||||
max_results = params.get("max_results", 50)
|
||||
|
||||
# Validate table
|
||||
valid_tables = ["files", "symbols", "metadata"]
|
||||
if table not in valid_tables:
|
||||
return f"Error: Invalid table '{table}'. Must be one of: {', '.join(valid_tables)}"
|
||||
|
||||
# Get project database
|
||||
try:
|
||||
if project_path:
|
||||
project_db = self.db_manager.get_project_db(project_path)
|
||||
else:
|
||||
import os
|
||||
|
||||
project_db = self.db_manager.get_project_for_path(os.getcwd())
|
||||
|
||||
if not project_db:
|
||||
return "Error: Could not find project database"
|
||||
|
||||
except PermissionError as e:
|
||||
return str(e)
|
||||
except Exception as e:
|
||||
return f"Error accessing project database: {str(e)}"
|
||||
|
||||
await tool_ctx.info(f"Searching in {table} table for pattern: {pattern}")
|
||||
|
||||
# Build search query
|
||||
conn = None
|
||||
try:
|
||||
conn = project_db.get_sqlite_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Get searchable columns for the table
|
||||
if column:
|
||||
# Validate column exists
|
||||
cursor.execute(f"PRAGMA table_info({table})")
|
||||
columns_info = cursor.fetchall()
|
||||
column_names = [col[1] for col in columns_info]
|
||||
|
||||
if column not in column_names:
|
||||
return f"Error: Column '{column}' not found in table '{table}'. Available columns: {', '.join(column_names)}"
|
||||
|
||||
search_columns = [column]
|
||||
else:
|
||||
# Get all text columns
|
||||
search_columns = self._get_text_columns(cursor, table)
|
||||
if not search_columns:
|
||||
return f"Error: No searchable text columns in table '{table}'"
|
||||
|
||||
# Build WHERE clause
|
||||
where_conditions = [f"{col} LIKE ?" for col in search_columns]
|
||||
where_clause = " OR ".join(where_conditions)
|
||||
|
||||
# Build query
|
||||
if table == "files":
|
||||
query = f"""
|
||||
SELECT path, SUBSTR(content, 1, 200) as snippet, size, modified_at
|
||||
FROM {table}
|
||||
WHERE {where_clause}
|
||||
LIMIT ?
|
||||
"""
|
||||
params_list = [pattern] * len(search_columns) + [max_results]
|
||||
|
||||
elif table == "symbols":
|
||||
query = f"""
|
||||
SELECT name, type, file_path, line_start, signature
|
||||
FROM {table}
|
||||
WHERE {where_clause}
|
||||
ORDER BY type, name
|
||||
LIMIT ?
|
||||
"""
|
||||
params_list = [pattern] * len(search_columns) + [max_results]
|
||||
|
||||
else: # metadata
|
||||
query = f"""
|
||||
SELECT key, value, updated_at
|
||||
FROM {table}
|
||||
WHERE {where_clause}
|
||||
LIMIT ?
|
||||
"""
|
||||
params_list = [pattern] * len(search_columns) + [max_results]
|
||||
|
||||
# Execute search
|
||||
cursor.execute(query, params_list)
|
||||
results = cursor.fetchall()
|
||||
|
||||
if not results:
|
||||
return f"No results found for pattern '{pattern}' in {table}"
|
||||
|
||||
# Format results
|
||||
output = self._format_results(table, results, pattern, search_columns)
|
||||
|
||||
return f"Found {len(results)} result(s) in {table}:\n\n{output}"
|
||||
|
||||
except sqlite3.Error as e:
|
||||
await tool_ctx.error(f"SQL error: {str(e)}")
|
||||
return f"SQL error: {str(e)}"
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Unexpected error: {str(e)}")
|
||||
return f"Error executing search: {str(e)}"
|
||||
finally:
|
||||
if conn:
|
||||
conn.close()
|
||||
|
||||
def _get_text_columns(self, cursor: sqlite3.Cursor, table: str) -> list[str]:
|
||||
"""Get text columns for a table."""
|
||||
cursor.execute(f"PRAGMA table_info({table})")
|
||||
columns_info = cursor.fetchall()
|
||||
|
||||
# Get TEXT columns
|
||||
text_columns = []
|
||||
for col in columns_info:
|
||||
col_name = col[1]
|
||||
col_type = col[2].upper()
|
||||
if "TEXT" in col_type or "CHAR" in col_type or col_type == "":
|
||||
text_columns.append(col_name)
|
||||
|
||||
return text_columns
|
||||
|
||||
def _format_results(
|
||||
self, table: str, results: list, pattern: str, search_columns: list[str]
|
||||
) -> str:
|
||||
"""Format search results based on table type."""
|
||||
output = []
|
||||
|
||||
if table == "files":
|
||||
output.append(f"Searched columns: {', '.join(search_columns)}\n")
|
||||
for row in results:
|
||||
path, snippet, size, modified = row
|
||||
output.append(f"File: {path}")
|
||||
output.append(f"Size: {size} bytes")
|
||||
output.append(f"Modified: {modified}")
|
||||
if snippet:
|
||||
# Highlight pattern in snippet
|
||||
snippet = snippet.replace("\n", " ")
|
||||
if len(snippet) > 150:
|
||||
snippet = snippet[:150] + "..."
|
||||
output.append(f"Content: {snippet}")
|
||||
output.append("-" * 60)
|
||||
|
||||
elif table == "symbols":
|
||||
output.append(f"Searched columns: {', '.join(search_columns)}\n")
|
||||
for row in results:
|
||||
name, type_, file_path, line_start, signature = row
|
||||
output.append(f"{type_}: {name}")
|
||||
output.append(f"File: {file_path}:{line_start}")
|
||||
if signature:
|
||||
output.append(f"Signature: {signature}")
|
||||
output.append("-" * 60)
|
||||
|
||||
else: # metadata
|
||||
output.append(f"Searched columns: {', '.join(search_columns)}\n")
|
||||
for row in results:
|
||||
key, value, updated = row
|
||||
output.append(f"Key: {key}")
|
||||
output.append(f"Value: {value}")
|
||||
output.append(f"Updated: {updated}")
|
||||
output.append("-" * 60)
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,266 +0,0 @@
|
||||
"""SQL statistics tool for database insights."""
|
||||
|
||||
import sqlite3
|
||||
from typing import Unpack, Optional, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.database.database_manager import DatabaseManager
|
||||
|
||||
ProjectPath = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Project path (defaults to current directory)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Detailed = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Show detailed statistics",
|
||||
default=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class SqlStatsParams(TypedDict, total=False):
|
||||
"""Parameters for SQL stats tool."""
|
||||
|
||||
project_path: Optional[str]
|
||||
detailed: bool
|
||||
|
||||
|
||||
@final
|
||||
class SqlStatsTool(BaseTool):
|
||||
"""Tool for getting SQLite database statistics."""
|
||||
|
||||
def __init__(
|
||||
self, permission_manager: PermissionManager, db_manager: DatabaseManager
|
||||
):
|
||||
"""Initialize the SQL stats tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
db_manager: Database manager instance
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.db_manager = db_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "sql_stats"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Get statistics about the project's SQLite database.
|
||||
|
||||
Shows:
|
||||
- Database size and location
|
||||
- Table information (row counts, sizes)
|
||||
- Index information
|
||||
- Column statistics
|
||||
- Most common values (with --detailed)
|
||||
|
||||
Examples:
|
||||
- sql_stats # Basic stats for current project
|
||||
- sql_stats --detailed # Detailed statistics
|
||||
- sql_stats --project-path /path/to/project
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[SqlStatsParams],
|
||||
) -> str:
|
||||
"""Get database statistics.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Database statistics
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
project_path = params.get("project_path")
|
||||
detailed = params.get("detailed", False)
|
||||
|
||||
# Get project database
|
||||
try:
|
||||
if project_path:
|
||||
project_db = self.db_manager.get_project_db(project_path)
|
||||
else:
|
||||
import os
|
||||
|
||||
project_db = self.db_manager.get_project_for_path(os.getcwd())
|
||||
|
||||
if not project_db:
|
||||
return "Error: Could not find project database"
|
||||
|
||||
except PermissionError as e:
|
||||
return str(e)
|
||||
except Exception as e:
|
||||
return f"Error accessing project database: {str(e)}"
|
||||
|
||||
await tool_ctx.info(
|
||||
f"Getting statistics for project: {project_db.project_path}"
|
||||
)
|
||||
|
||||
# Collect statistics
|
||||
conn = None
|
||||
try:
|
||||
conn = project_db.get_sqlite_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
output = []
|
||||
output.append(f"=== SQLite Database Statistics ===")
|
||||
output.append(f"Project: {project_db.project_path}")
|
||||
output.append(f"Database: {project_db.sqlite_path}")
|
||||
|
||||
# Get database size
|
||||
db_size = project_db.sqlite_path.stat().st_size
|
||||
output.append(f"Database Size: {self._format_size(db_size)}")
|
||||
output.append("")
|
||||
|
||||
# Get table statistics
|
||||
cursor.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"
|
||||
)
|
||||
tables = cursor.fetchall()
|
||||
|
||||
output.append("=== Tables ===")
|
||||
total_rows = 0
|
||||
|
||||
for (table_name,) in tables:
|
||||
if table_name.startswith("sqlite_"):
|
||||
continue
|
||||
|
||||
# Get row count
|
||||
cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
|
||||
row_count = cursor.fetchone()[0]
|
||||
total_rows += row_count
|
||||
|
||||
# Get table info
|
||||
cursor.execute(f"PRAGMA table_info({table_name})")
|
||||
columns = cursor.fetchall()
|
||||
col_count = len(columns)
|
||||
|
||||
output.append(f"\n{table_name}:")
|
||||
output.append(f" Rows: {row_count:,}")
|
||||
output.append(f" Columns: {col_count}")
|
||||
|
||||
if detailed and row_count > 0:
|
||||
# Show column details
|
||||
output.append(" Columns:")
|
||||
for col in columns:
|
||||
col_name = col[1]
|
||||
col_type = col[2]
|
||||
is_pk = col[5]
|
||||
not_null = col[3]
|
||||
|
||||
flags = []
|
||||
if is_pk:
|
||||
flags.append("PRIMARY KEY")
|
||||
if not_null:
|
||||
flags.append("NOT NULL")
|
||||
|
||||
flag_str = f" ({', '.join(flags)})" if flags else ""
|
||||
output.append(f" - {col_name}: {col_type}{flag_str}")
|
||||
|
||||
# Show sample data for specific tables
|
||||
if table_name == "files" and row_count > 0:
|
||||
cursor.execute(
|
||||
f"SELECT COUNT(DISTINCT SUBSTR(path, -3)) as ext_count FROM {table_name}"
|
||||
)
|
||||
ext_count = cursor.fetchone()[0]
|
||||
output.append(f" File types: ~{ext_count}")
|
||||
|
||||
elif table_name == "symbols" and row_count > 0:
|
||||
cursor.execute(
|
||||
f"SELECT type, COUNT(*) as count FROM {table_name} GROUP BY type ORDER BY count DESC LIMIT 5"
|
||||
)
|
||||
symbol_types = cursor.fetchall()
|
||||
output.append(" Symbol types:")
|
||||
for sym_type, count in symbol_types:
|
||||
output.append(f" - {sym_type}: {count}")
|
||||
|
||||
# Get indexes
|
||||
cursor.execute(f"PRAGMA index_list({table_name})")
|
||||
indexes = cursor.fetchall()
|
||||
if indexes:
|
||||
output.append(f" Indexes: {len(indexes)}")
|
||||
|
||||
output.append(f"\nTotal Rows: {total_rows:,}")
|
||||
|
||||
# Get index statistics
|
||||
cursor.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='index' AND sql IS NOT NULL ORDER BY name"
|
||||
)
|
||||
indexes = cursor.fetchall()
|
||||
if indexes:
|
||||
output.append(f"\n=== Indexes ===")
|
||||
output.append(f"Total Indexes: {len(indexes)}")
|
||||
|
||||
if detailed:
|
||||
for (idx_name,) in indexes:
|
||||
cursor.execute(f"PRAGMA index_info({idx_name})")
|
||||
idx_info = cursor.fetchall()
|
||||
if idx_info:
|
||||
cols = [info[2] for info in idx_info]
|
||||
output.append(f" {idx_name}: ({', '.join(cols)})")
|
||||
|
||||
# Database properties
|
||||
if detailed:
|
||||
output.append("\n=== Database Properties ===")
|
||||
|
||||
# Page size
|
||||
cursor.execute("PRAGMA page_size")
|
||||
page_size = cursor.fetchone()[0]
|
||||
output.append(f"Page Size: {page_size:,} bytes")
|
||||
|
||||
# Page count
|
||||
cursor.execute("PRAGMA page_count")
|
||||
page_count = cursor.fetchone()[0]
|
||||
output.append(f"Page Count: {page_count:,}")
|
||||
|
||||
# Cache size
|
||||
cursor.execute("PRAGMA cache_size")
|
||||
cache_size = cursor.fetchone()[0]
|
||||
output.append(f"Cache Size: {abs(cache_size):,} pages")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except sqlite3.Error as e:
|
||||
await tool_ctx.error(f"SQL error: {str(e)}")
|
||||
return f"SQL error: {str(e)}"
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Unexpected error: {str(e)}")
|
||||
return f"Error getting statistics: {str(e)}"
|
||||
finally:
|
||||
if conn:
|
||||
conn.close()
|
||||
|
||||
def _format_size(self, size: int) -> str:
|
||||
"""Format file size in human-readable format."""
|
||||
for unit in ["B", "KB", "MB", "GB"]:
|
||||
if size < 1024.0:
|
||||
return f"{size:.1f} {unit}"
|
||||
size /= 1024.0
|
||||
return f"{size:.1f} TB"
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,11 +0,0 @@
|
||||
"""Editor integration tools for Hanzo AI."""
|
||||
|
||||
from hanzo_mcp.tools.editor.neovim_edit import NeovimEditTool
|
||||
from hanzo_mcp.tools.editor.neovim_command import NeovimCommandTool
|
||||
from hanzo_mcp.tools.editor.neovim_session import NeovimSessionTool
|
||||
|
||||
__all__ = [
|
||||
"NeovimEditTool",
|
||||
"NeovimCommandTool",
|
||||
"NeovimSessionTool",
|
||||
]
|
||||
@@ -1,269 +0,0 @@
|
||||
"""Execute Neovim commands and macros."""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import subprocess
|
||||
from typing import List, Unpack, Optional, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
Command = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Neovim command to execute (Ex commands like :w, :q, etc.)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Commands = Annotated[
|
||||
Optional[List[str]],
|
||||
Field(
|
||||
description="List of Neovim commands to execute in sequence",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Macro = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Vim macro to execute (e.g., 'dd' to delete line)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
FilePath = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="File to operate on (optional, uses current buffer if not specified)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
SaveAfter = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Save file after executing commands",
|
||||
default=True,
|
||||
),
|
||||
]
|
||||
|
||||
ReturnOutput = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Return output/messages from Neovim",
|
||||
default=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class NeovimCommandParams(TypedDict, total=False):
|
||||
"""Parameters for Neovim command tool."""
|
||||
|
||||
command: Optional[str]
|
||||
commands: Optional[List[str]]
|
||||
macro: Optional[str]
|
||||
file_path: Optional[str]
|
||||
save_after: bool
|
||||
return_output: bool
|
||||
|
||||
|
||||
@final
|
||||
class NeovimCommandTool(BaseTool):
|
||||
"""Tool for executing Neovim commands and macros."""
|
||||
|
||||
def __init__(self, permission_manager: PermissionManager):
|
||||
"""Initialize the Neovim command tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "neovim_command"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Execute Neovim commands and macros programmatically.
|
||||
|
||||
Run Ex commands, normal mode commands, or complex macros in Neovim.
|
||||
Can operate on files without opening the editor interface.
|
||||
|
||||
Examples:
|
||||
- neovim_command --command ":set number" --file-path main.py
|
||||
- neovim_command --command ":%s/old/new/g" --file-path config.json
|
||||
- neovim_command --commands ":set expandtab" ":retab" --file-path script.sh
|
||||
- neovim_command --macro "ggVG=" --file-path messy.py # Format entire file
|
||||
- neovim_command --macro "dd10j" --file-path list.txt # Delete line and go down 10
|
||||
|
||||
Common commands:
|
||||
- :w - Save file
|
||||
- :q - Quit
|
||||
- :%s/old/new/g - Replace all occurrences
|
||||
- :set number - Show line numbers
|
||||
- :set expandtab - Use spaces instead of tabs
|
||||
- :retab - Convert tabs to spaces
|
||||
|
||||
Common macros:
|
||||
- gg - Go to beginning of file
|
||||
- G - Go to end of file
|
||||
- dd - Delete line
|
||||
- yy - Yank (copy) line
|
||||
- p - Paste
|
||||
- V - Visual line mode
|
||||
- = - Format/indent
|
||||
|
||||
Note: Requires Neovim to be installed.
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[NeovimCommandParams],
|
||||
) -> str:
|
||||
"""Execute Neovim command.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Result of the command execution
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
command = params.get("command")
|
||||
commands = params.get("commands")
|
||||
macro = params.get("macro")
|
||||
file_path = params.get("file_path")
|
||||
save_after = params.get("save_after", True)
|
||||
return_output = params.get("return_output", True)
|
||||
|
||||
# Validate inputs
|
||||
if not any([command, commands, macro]):
|
||||
return "Error: Must provide either 'command', 'commands', or 'macro'"
|
||||
|
||||
if sum(bool(x) for x in [command, commands, macro]) > 1:
|
||||
return (
|
||||
"Error: Can only use one of 'command', 'commands', or 'macro' at a time"
|
||||
)
|
||||
|
||||
# Check if Neovim is available
|
||||
nvim_cmd = shutil.which("nvim")
|
||||
if not nvim_cmd:
|
||||
return "Error: Neovim (nvim) not found. Install it first."
|
||||
|
||||
# Prepare commands list
|
||||
nvim_commands = []
|
||||
|
||||
if command:
|
||||
nvim_commands.append(command)
|
||||
elif commands:
|
||||
nvim_commands.extend(commands)
|
||||
elif macro:
|
||||
# Convert macro to normal mode command
|
||||
# Escape special characters
|
||||
escaped_macro = macro.replace('"', '\\"')
|
||||
nvim_commands.append(f':normal "{escaped_macro}"')
|
||||
|
||||
# Add save command if requested
|
||||
if save_after:
|
||||
nvim_commands.append(":w")
|
||||
|
||||
# Always quit at the end
|
||||
nvim_commands.append(":q")
|
||||
|
||||
# Build Neovim command line
|
||||
cmd = [nvim_cmd, "-n", "-i", "NONE"] # No swap file, no shada file
|
||||
|
||||
# Add commands
|
||||
for vim_cmd in nvim_commands:
|
||||
cmd.extend(["-c", vim_cmd])
|
||||
|
||||
# Add file if specified
|
||||
if file_path:
|
||||
file_path = os.path.abspath(file_path)
|
||||
|
||||
# Check permissions
|
||||
if not self.permission_manager.has_permission(file_path):
|
||||
return f"Error: No permission to access {file_path}"
|
||||
|
||||
if not os.path.exists(file_path):
|
||||
return f"Error: File not found: {file_path}"
|
||||
|
||||
cmd.append(file_path)
|
||||
else:
|
||||
# Create empty buffer
|
||||
cmd.append("-")
|
||||
|
||||
await tool_ctx.info(f"Executing Neovim commands: {nvim_commands}")
|
||||
|
||||
try:
|
||||
# Execute Neovim
|
||||
if return_output:
|
||||
# Capture output by redirecting messages
|
||||
output_file = tempfile.NamedTemporaryFile(mode="w+", delete=False)
|
||||
output_file.close()
|
||||
|
||||
# Add command to redirect messages
|
||||
cmd.insert(3, "-c")
|
||||
cmd.insert(4, f":redir! > {output_file.name}")
|
||||
|
||||
# Execute
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
|
||||
# Read output
|
||||
output_content = ""
|
||||
try:
|
||||
with open(output_file.name, "r") as f:
|
||||
output_content = f.read().strip()
|
||||
finally:
|
||||
os.unlink(output_file.name)
|
||||
|
||||
if result.returncode == 0:
|
||||
response = "Commands executed successfully"
|
||||
if file_path:
|
||||
response += f" on {os.path.basename(file_path)}"
|
||||
if output_content:
|
||||
response += f"\n\nOutput:\n{output_content}"
|
||||
return response
|
||||
else:
|
||||
error_msg = "Error executing Neovim commands"
|
||||
if result.stderr:
|
||||
error_msg += f"\n\nError:\n{result.stderr}"
|
||||
if output_content:
|
||||
error_msg += f"\n\nOutput:\n{output_content}"
|
||||
return error_msg
|
||||
else:
|
||||
# Just execute without capturing output
|
||||
result = subprocess.run(cmd)
|
||||
|
||||
if result.returncode == 0:
|
||||
response = "Commands executed successfully"
|
||||
if file_path:
|
||||
response += f" on {os.path.basename(file_path)}"
|
||||
return response
|
||||
else:
|
||||
return f"Neovim exited with code {result.returncode}"
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to execute Neovim commands: {str(e)}")
|
||||
return f"Error executing Neovim commands: {str(e)}"
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,288 +0,0 @@
|
||||
"""Open files in Neovim editor."""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from typing import Unpack, Optional, Annotated, TypedDict, final, override
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
FilePath = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Path to the file to open",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
LineNumber = Annotated[
|
||||
Optional[int],
|
||||
Field(
|
||||
description="Line number to jump to",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
ColumnNumber = Annotated[
|
||||
Optional[int],
|
||||
Field(
|
||||
description="Column number to jump to",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
ReadOnly = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Open file in read-only mode",
|
||||
default=False,
|
||||
),
|
||||
]
|
||||
|
||||
Split = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Split mode: vsplit, split, tab",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Wait = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Wait for Neovim to exit before returning",
|
||||
default=True,
|
||||
),
|
||||
]
|
||||
|
||||
InTerminal = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Open in terminal (requires terminal that supports it)",
|
||||
default=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class NeovimEditParams(TypedDict, total=False):
|
||||
"""Parameters for Neovim edit tool."""
|
||||
|
||||
file_path: str
|
||||
line_number: Optional[int]
|
||||
column_number: Optional[int]
|
||||
read_only: bool
|
||||
split: Optional[str]
|
||||
wait: bool
|
||||
in_terminal: bool
|
||||
|
||||
|
||||
@final
|
||||
class NeovimEditTool(BaseTool):
|
||||
"""Tool for opening files in Neovim."""
|
||||
|
||||
def __init__(self, permission_manager: PermissionManager):
|
||||
"""Initialize the Neovim edit tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "neovim_edit"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Open files in Neovim editor with advanced options.
|
||||
|
||||
Open files at specific lines/columns, in different split modes, or read-only.
|
||||
Integrates with your existing Neovim configuration.
|
||||
|
||||
Examples:
|
||||
- neovim_edit --file-path main.py
|
||||
- neovim_edit --file-path main.py --line-number 42
|
||||
- neovim_edit --file-path main.py --line-number 42 --column-number 10
|
||||
- neovim_edit --file-path config.json --read-only
|
||||
- neovim_edit --file-path test.py --split vsplit
|
||||
- neovim_edit --file-path README.md --split tab
|
||||
|
||||
Split modes:
|
||||
- vsplit: Open in vertical split
|
||||
- split: Open in horizontal split
|
||||
- tab: Open in new tab
|
||||
|
||||
Note: Requires Neovim to be installed and available in PATH.
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[NeovimEditParams],
|
||||
) -> str:
|
||||
"""Open file in Neovim.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Result of the operation
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
file_path = params.get("file_path")
|
||||
if not file_path:
|
||||
return "Error: file_path is required"
|
||||
|
||||
line_number = params.get("line_number")
|
||||
column_number = params.get("column_number")
|
||||
read_only = params.get("read_only", False)
|
||||
split = params.get("split")
|
||||
wait = params.get("wait", True)
|
||||
in_terminal = params.get("in_terminal", True)
|
||||
|
||||
# Check if Neovim is available
|
||||
nvim_cmd = shutil.which("nvim")
|
||||
if not nvim_cmd:
|
||||
# Try common locations
|
||||
common_paths = [
|
||||
"/usr/local/bin/nvim",
|
||||
"/usr/bin/nvim",
|
||||
"/opt/homebrew/bin/nvim",
|
||||
os.path.expanduser("~/.local/bin/nvim"),
|
||||
]
|
||||
for path in common_paths:
|
||||
if os.path.exists(path):
|
||||
nvim_cmd = path
|
||||
break
|
||||
|
||||
if not nvim_cmd:
|
||||
return """Error: Neovim (nvim) not found. Install it with:
|
||||
|
||||
On macOS:
|
||||
brew install neovim
|
||||
|
||||
On Ubuntu/Debian:
|
||||
sudo apt install neovim
|
||||
|
||||
On Arch:
|
||||
sudo pacman -S neovim
|
||||
|
||||
Or visit: https://neovim.io/"""
|
||||
|
||||
# Convert to absolute path
|
||||
file_path = os.path.abspath(file_path)
|
||||
|
||||
# Check permissions
|
||||
if not self.permission_manager.has_permission(file_path):
|
||||
return f"Error: No permission to access {file_path}"
|
||||
|
||||
# Build Neovim command
|
||||
cmd = [nvim_cmd]
|
||||
|
||||
# Add read-only flag
|
||||
if read_only:
|
||||
cmd.append("-R")
|
||||
|
||||
# Add split mode
|
||||
if split:
|
||||
if split == "vsplit":
|
||||
cmd.extend(["-c", "vsplit"])
|
||||
elif split == "split":
|
||||
cmd.extend(["-c", "split"])
|
||||
elif split == "tab":
|
||||
cmd.extend(["-c", "tabnew"])
|
||||
else:
|
||||
return f"Error: Invalid split mode '{split}'. Use 'vsplit', 'split', or 'tab'"
|
||||
|
||||
# Add file path
|
||||
cmd.append(file_path)
|
||||
|
||||
# Add line/column positioning
|
||||
if line_number:
|
||||
if column_number:
|
||||
# Go to specific line and column
|
||||
cmd.extend(["+call cursor({}, {})".format(line_number, column_number)])
|
||||
else:
|
||||
# Go to specific line
|
||||
cmd.append(f"+{line_number}")
|
||||
|
||||
await tool_ctx.info(f"Opening {file_path} in Neovim")
|
||||
|
||||
try:
|
||||
# Determine how to run Neovim
|
||||
if in_terminal and not wait:
|
||||
# Open in a new terminal window (platform-specific)
|
||||
if os.uname().sysname == "Darwin": # macOS
|
||||
# Try to use iTerm2 if available, otherwise Terminal
|
||||
if shutil.which("osascript"):
|
||||
# Build AppleScript to open in iTerm2 or Terminal
|
||||
nvim_cmd_str = " ".join(f'"{arg}"' for arg in cmd)
|
||||
|
||||
# Try iTerm2 first
|
||||
applescript = f"""tell application "System Events"
|
||||
if exists application process "iTerm2" then
|
||||
tell application "iTerm"
|
||||
activate
|
||||
tell current window
|
||||
create tab with default profile
|
||||
tell current session
|
||||
write text "{nvim_cmd_str}"
|
||||
end tell
|
||||
end tell
|
||||
end tell
|
||||
else
|
||||
tell application "Terminal"
|
||||
activate
|
||||
do script "{nvim_cmd_str}"
|
||||
end tell
|
||||
end if
|
||||
end tell"""
|
||||
|
||||
subprocess.run(["osascript", "-e", applescript])
|
||||
return f"Opened {file_path} in Neovim (new terminal window)"
|
||||
|
||||
elif shutil.which("gnome-terminal"):
|
||||
# Linux with GNOME
|
||||
subprocess.Popen(["gnome-terminal", "--"] + cmd)
|
||||
return f"Opened {file_path} in Neovim (new terminal window)"
|
||||
|
||||
elif shutil.which("xterm"):
|
||||
# Fallback to xterm
|
||||
subprocess.Popen(["xterm", "-e"] + cmd)
|
||||
return f"Opened {file_path} in Neovim (new terminal window)"
|
||||
|
||||
else:
|
||||
# Can't open in terminal, fall back to subprocess
|
||||
subprocess.Popen(cmd)
|
||||
return f"Opened {file_path} in Neovim (background process)"
|
||||
|
||||
else:
|
||||
# Run and wait for completion
|
||||
result = subprocess.run(cmd)
|
||||
|
||||
if result.returncode == 0:
|
||||
return f"Successfully edited {file_path} in Neovim"
|
||||
else:
|
||||
return f"Neovim exited with code {result.returncode}"
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to open Neovim: {str(e)}")
|
||||
return f"Error opening Neovim: {str(e)}"
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,362 +0,0 @@
|
||||
"""Manage Neovim sessions."""
|
||||
|
||||
import os
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
from typing import Unpack, Optional, Annotated, TypedDict, final, override
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
|
||||
Action = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Action to perform: save, restore, list, delete",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
SessionName = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Name of the session",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
ProjectPath = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Project path (defaults to current directory)",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
AutoName = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Auto-generate session name based on project and timestamp",
|
||||
default=False,
|
||||
),
|
||||
]
|
||||
|
||||
Overwrite = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Overwrite existing session",
|
||||
default=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class NeovimSessionParams(TypedDict, total=False):
|
||||
"""Parameters for Neovim session tool."""
|
||||
|
||||
action: str
|
||||
session_name: Optional[str]
|
||||
project_path: Optional[str]
|
||||
auto_name: bool
|
||||
overwrite: bool
|
||||
|
||||
|
||||
@final
|
||||
class NeovimSessionTool(BaseTool):
|
||||
"""Tool for managing Neovim sessions."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the Neovim session tool."""
|
||||
self.session_dir = Path.home() / ".hanzo" / "neovim" / "sessions"
|
||||
self.session_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "neovim_session"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Save and restore Neovim editing sessions.
|
||||
|
||||
Manage Neovim sessions to save your workspace state including:
|
||||
- Open files and buffers
|
||||
- Window layouts and splits
|
||||
- Cursor positions
|
||||
- Marks and registers
|
||||
- Local options and mappings
|
||||
|
||||
Actions:
|
||||
- save: Save current Neovim session
|
||||
- restore: Restore a saved session
|
||||
- list: List all saved sessions
|
||||
- delete: Delete a saved session
|
||||
|
||||
Examples:
|
||||
- neovim_session --action save --session-name "feature-work"
|
||||
- neovim_session --action save --auto-name # Auto-generate name
|
||||
- neovim_session --action restore --session-name "feature-work"
|
||||
- neovim_session --action list
|
||||
- neovim_session --action list --project-path /path/to/project
|
||||
- neovim_session --action delete --session-name "old-session"
|
||||
|
||||
Sessions are stored in ~/.hanzo/neovim/sessions/
|
||||
Project-specific sessions are automatically organized by project path.
|
||||
|
||||
Note: Requires Neovim to be installed.
|
||||
"""
|
||||
|
||||
@override
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[NeovimSessionParams],
|
||||
) -> str:
|
||||
"""Manage Neovim session.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Result of the session operation
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
action = params.get("action")
|
||||
if not action:
|
||||
return "Error: action is required (save, restore, list, delete)"
|
||||
|
||||
session_name = params.get("session_name")
|
||||
project_path = params.get("project_path") or os.getcwd()
|
||||
auto_name = params.get("auto_name", False)
|
||||
overwrite = params.get("overwrite", False)
|
||||
|
||||
# Validate action
|
||||
valid_actions = ["save", "restore", "list", "delete"]
|
||||
if action not in valid_actions:
|
||||
return f"Error: Invalid action '{action}'. Must be one of: {', '.join(valid_actions)}"
|
||||
|
||||
# Check if Neovim is available
|
||||
nvim_cmd = shutil.which("nvim")
|
||||
if not nvim_cmd and action in ["save", "restore"]:
|
||||
return "Error: Neovim (nvim) not found. Install it first."
|
||||
|
||||
# Get project-specific session directory
|
||||
project_hash = str(hash(os.path.abspath(project_path)) % 10**8)
|
||||
project_name = os.path.basename(project_path) or "root"
|
||||
project_session_dir = self.session_dir / f"{project_name}_{project_hash}"
|
||||
project_session_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Handle different actions
|
||||
if action == "save":
|
||||
return await self._save_session(
|
||||
tool_ctx,
|
||||
session_name,
|
||||
project_session_dir,
|
||||
auto_name,
|
||||
overwrite,
|
||||
project_path,
|
||||
)
|
||||
elif action == "restore":
|
||||
return await self._restore_session(
|
||||
tool_ctx, session_name, project_session_dir
|
||||
)
|
||||
elif action == "list":
|
||||
return self._list_sessions(project_session_dir, project_path)
|
||||
elif action == "delete":
|
||||
return self._delete_session(session_name, project_session_dir)
|
||||
|
||||
async def _save_session(
|
||||
self,
|
||||
tool_ctx,
|
||||
session_name: Optional[str],
|
||||
project_dir: Path,
|
||||
auto_name: bool,
|
||||
overwrite: bool,
|
||||
project_path: str,
|
||||
) -> str:
|
||||
"""Save Neovim session."""
|
||||
# Generate session name if needed
|
||||
if auto_name or not session_name:
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
session_name = f"session_{timestamp}"
|
||||
|
||||
# Sanitize session name
|
||||
session_name = session_name.replace("/", "_").replace(" ", "_")
|
||||
|
||||
session_file = project_dir / f"{session_name}.vim"
|
||||
metadata_file = project_dir / f"{session_name}.json"
|
||||
|
||||
# Check if exists
|
||||
if session_file.exists() and not overwrite:
|
||||
return f"Error: Session '{session_name}' already exists. Use --overwrite to replace."
|
||||
|
||||
await tool_ctx.info(f"Saving Neovim session: {session_name}")
|
||||
|
||||
# Create temporary vim script to save session
|
||||
vim_script = f"""
|
||||
:mksession! {session_file}
|
||||
:echo "Session saved to {session_file}"
|
||||
:qa
|
||||
"""
|
||||
|
||||
try:
|
||||
# Run Neovim to save session
|
||||
# First, check if Neovim is already running
|
||||
# For now, we'll create a new instance
|
||||
result = subprocess.run(
|
||||
["nvim", "-c", vim_script.strip()], capture_output=True, text=True
|
||||
)
|
||||
|
||||
if result.returncode != 0 and result.stderr:
|
||||
return f"Error saving session: {result.stderr}"
|
||||
|
||||
# Save metadata
|
||||
metadata = {
|
||||
"name": session_name,
|
||||
"created_at": datetime.now().isoformat(),
|
||||
"project_path": project_path,
|
||||
"description": f"Neovim session for {project_path}",
|
||||
}
|
||||
|
||||
with open(metadata_file, "w") as f:
|
||||
json.dump(metadata, f, indent=2)
|
||||
|
||||
return f"""Successfully saved Neovim session '{session_name}'
|
||||
|
||||
Session file: {session_file}
|
||||
Project: {project_path}
|
||||
|
||||
To restore this session:
|
||||
neovim_session --action restore --session-name "{session_name}"
|
||||
|
||||
Or manually in Neovim:
|
||||
:source {session_file}"""
|
||||
|
||||
except Exception as e:
|
||||
return f"Error saving session: {str(e)}"
|
||||
|
||||
async def _restore_session(
|
||||
self, tool_ctx, session_name: Optional[str], project_dir: Path
|
||||
) -> str:
|
||||
"""Restore Neovim session."""
|
||||
if not session_name:
|
||||
# List available sessions
|
||||
sessions = list(project_dir.glob("*.vim"))
|
||||
if not sessions:
|
||||
return "Error: No sessions found for this project. Use 'neovim_session --action list' to see all sessions."
|
||||
|
||||
# Use most recent
|
||||
sessions.sort(key=lambda x: x.stat().st_mtime, reverse=True)
|
||||
session_file = sessions[0]
|
||||
session_name = session_file.stem
|
||||
else:
|
||||
session_file = project_dir / f"{session_name}.vim"
|
||||
if not session_file.exists():
|
||||
return f"Error: Session '{session_name}' not found. Use 'neovim_session --action list' to see available sessions."
|
||||
|
||||
await tool_ctx.info(f"Restoring Neovim session: {session_name}")
|
||||
|
||||
try:
|
||||
# Open Neovim with the session
|
||||
subprocess.run(["nvim", "-S", str(session_file)])
|
||||
|
||||
return f"Restored Neovim session '{session_name}'"
|
||||
|
||||
except Exception as e:
|
||||
return f"Error restoring session: {str(e)}"
|
||||
|
||||
def _list_sessions(self, project_dir: Path, project_path: str) -> str:
|
||||
"""List available sessions."""
|
||||
output = ["=== Neovim Sessions ==="]
|
||||
output.append(f"Project: {project_path}\n")
|
||||
|
||||
# List project-specific sessions
|
||||
sessions = list(project_dir.glob("*.vim"))
|
||||
|
||||
if sessions:
|
||||
output.append("Project Sessions:")
|
||||
sessions.sort(key=lambda x: x.stat().st_mtime, reverse=True)
|
||||
|
||||
for session_file in sessions:
|
||||
session_name = session_file.stem
|
||||
metadata_file = project_dir / f"{session_name}.json"
|
||||
|
||||
# Get metadata if available
|
||||
created_at = "Unknown"
|
||||
if metadata_file.exists():
|
||||
try:
|
||||
with open(metadata_file, "r") as f:
|
||||
metadata = json.load(f)
|
||||
created_at = metadata.get("created_at", "Unknown")
|
||||
if created_at != "Unknown":
|
||||
# Format date
|
||||
dt = datetime.fromisoformat(created_at)
|
||||
created_at = dt.strftime("%Y-%m-%d %H:%M:%S")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Get file size
|
||||
size = session_file.stat().st_size
|
||||
size_kb = size / 1024
|
||||
|
||||
output.append(f" - {session_name}")
|
||||
output.append(f" Created: {created_at}")
|
||||
output.append(f" Size: {size_kb:.1f} KB")
|
||||
else:
|
||||
output.append("No sessions found for this project.")
|
||||
|
||||
# Also list all sessions
|
||||
all_sessions = list(self.session_dir.rglob("*.vim"))
|
||||
other_sessions = [s for s in all_sessions if s.parent != project_dir]
|
||||
|
||||
if other_sessions:
|
||||
output.append("\nOther Projects' Sessions:")
|
||||
for session_file in other_sessions[:10]: # Show max 10
|
||||
project_name = session_file.parent.name
|
||||
session_name = session_file.stem
|
||||
output.append(f" - {project_name}/{session_name}")
|
||||
|
||||
if len(other_sessions) > 10:
|
||||
output.append(f" ... and {len(other_sessions) - 10} more")
|
||||
|
||||
output.append(
|
||||
"\nUse 'neovim_session --action restore --session-name <name>' to restore a session."
|
||||
)
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _delete_session(self, session_name: Optional[str], project_dir: Path) -> str:
|
||||
"""Delete a session."""
|
||||
if not session_name:
|
||||
return "Error: session_name is required for delete action"
|
||||
|
||||
session_file = project_dir / f"{session_name}.vim"
|
||||
metadata_file = project_dir / f"{session_name}.json"
|
||||
|
||||
if not session_file.exists():
|
||||
return f"Error: Session '{session_name}' not found"
|
||||
|
||||
try:
|
||||
session_file.unlink()
|
||||
if metadata_file.exists():
|
||||
metadata_file.unlink()
|
||||
|
||||
return f"Successfully deleted session '{session_name}'"
|
||||
|
||||
except Exception as e:
|
||||
return f"Error deleting session: {str(e)}"
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user