mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
Fix hanging
This commit is contained in:
@@ -129,29 +129,20 @@ batch(
|
||||
3. Check `has_next` before requesting next page
|
||||
4. Use batch for parallel page fetching"""
|
||||
|
||||
MEMORY_VECTOR_HELP_PROMPT = """# Memory & Vector Tools Guide
|
||||
MEMORY_VECTOR_HELP_PROMPT = """# Memory & Search Tools Guide
|
||||
|
||||
## Semantic Search and Memory
|
||||
## Search
|
||||
|
||||
### Vector Search
|
||||
Search using semantic similarity rather than exact text matching.
|
||||
### Unified Search
|
||||
Fast search using text, AST, and symbol matching.
|
||||
|
||||
```python
|
||||
# Enable vector search in unified search
|
||||
# Text and code pattern search
|
||||
search(
|
||||
pattern="how to handle user authentication",
|
||||
pattern="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
|
||||
enable_ast=true,
|
||||
enable_symbol=true
|
||||
)
|
||||
```
|
||||
|
||||
|
||||
@@ -222,7 +222,6 @@ grep_ast(pattern="class.*Controller", path="/project/src")
|
||||
search_content(
|
||||
pattern="authentication",
|
||||
path="/project",
|
||||
enable_vector=True,
|
||||
enable_ast=True
|
||||
)
|
||||
```
|
||||
|
||||
@@ -1,4 +1,10 @@
|
||||
"""MCP server implementing Hanzo capabilities."""
|
||||
"""MCP server implementing Hanzo capabilities.
|
||||
|
||||
IMPORTANT: This module uses lazy imports to ensure fast startup.
|
||||
Heavy imports happen only when the server is actually created, not at module load time.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import atexit
|
||||
@@ -7,26 +13,87 @@ import logging
|
||||
import secrets
|
||||
import warnings
|
||||
import threading
|
||||
from typing import Literal, cast, final
|
||||
from typing import TYPE_CHECKING, Literal, cast, final
|
||||
|
||||
# Type-only imports - don't execute at runtime
|
||||
if TYPE_CHECKING:
|
||||
from mcp.server import FastMCP
|
||||
from hanzo_mcp.server_enhanced import EnhancedFastMCP
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
from hanzo_mcp.tools.shell.session_storage import SessionStorage
|
||||
|
||||
# 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
|
||||
try:
|
||||
from mcp.server import FastMCP
|
||||
except ImportError:
|
||||
from mcp import FastMCP
|
||||
# Cached imports - lazy loaded on first use
|
||||
_FastMCP = None
|
||||
_EnhancedFastMCP = None
|
||||
_PermissionManager = None
|
||||
_SessionStorage = None
|
||||
_register_all_tools = None
|
||||
_register_all_prompts = None
|
||||
|
||||
# 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
|
||||
|
||||
def _get_fast_mcp():
|
||||
"""Get FastMCP class lazily."""
|
||||
global _FastMCP
|
||||
if _FastMCP is None:
|
||||
try:
|
||||
from fastmcp import FastMCP
|
||||
_FastMCP = FastMCP
|
||||
except ImportError:
|
||||
try:
|
||||
from mcp.server import FastMCP
|
||||
_FastMCP = FastMCP
|
||||
except ImportError:
|
||||
from mcp import FastMCP
|
||||
_FastMCP = FastMCP
|
||||
return _FastMCP
|
||||
|
||||
|
||||
def _get_enhanced_fast_mcp():
|
||||
"""Get EnhancedFastMCP class lazily."""
|
||||
global _EnhancedFastMCP
|
||||
if _EnhancedFastMCP is None:
|
||||
from hanzo_mcp.server_enhanced import EnhancedFastMCP
|
||||
_EnhancedFastMCP = EnhancedFastMCP
|
||||
return _EnhancedFastMCP
|
||||
|
||||
|
||||
def _get_permission_manager():
|
||||
"""Get PermissionManager class lazily."""
|
||||
global _PermissionManager
|
||||
if _PermissionManager is None:
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
_PermissionManager = PermissionManager
|
||||
return _PermissionManager
|
||||
|
||||
|
||||
def _get_session_storage():
|
||||
"""Get SessionStorage class lazily."""
|
||||
global _SessionStorage
|
||||
if _SessionStorage is None:
|
||||
from hanzo_mcp.tools.shell.session_storage import SessionStorage
|
||||
_SessionStorage = SessionStorage
|
||||
return _SessionStorage
|
||||
|
||||
|
||||
def _get_register_all_tools():
|
||||
"""Get register_all_tools function lazily."""
|
||||
global _register_all_tools
|
||||
if _register_all_tools is None:
|
||||
from hanzo_mcp.tools import register_all_tools
|
||||
_register_all_tools = register_all_tools
|
||||
return _register_all_tools
|
||||
|
||||
|
||||
def _get_register_all_prompts():
|
||||
"""Get register_all_prompts function lazily."""
|
||||
global _register_all_prompts
|
||||
if _register_all_prompts is None:
|
||||
from hanzo_mcp.prompts import register_all_prompts
|
||||
_register_all_prompts = register_all_prompts
|
||||
return _register_all_prompts
|
||||
|
||||
|
||||
@final
|
||||
@@ -80,6 +147,7 @@ class HanzoMCPServer:
|
||||
disabled_tools: List of tool names to disable (default: None)
|
||||
"""
|
||||
# Use enhanced server for automatic context normalization
|
||||
EnhancedFastMCP = _get_enhanced_fast_mcp()
|
||||
self.mcp = mcp_instance if mcp_instance is not None else EnhancedFastMCP(name)
|
||||
|
||||
# Initialize authentication token
|
||||
@@ -92,6 +160,7 @@ class HanzoMCPServer:
|
||||
logger.warning("Set HANZO_MCP_TOKEN environment variable for persistent auth")
|
||||
|
||||
# Initialize permissions and command executor
|
||||
PermissionManager = _get_permission_manager()
|
||||
self.permission_manager = PermissionManager()
|
||||
|
||||
# Handle project_dir parameter
|
||||
@@ -143,7 +212,8 @@ class HanzoMCPServer:
|
||||
# Store the final processed tool configuration
|
||||
self.enabled_tools = final_enabled_tools
|
||||
|
||||
# Register all tools
|
||||
# Register all tools (lazy import)
|
||||
register_all_tools = _get_register_all_tools()
|
||||
register_all_tools(
|
||||
mcp_server=self.mcp,
|
||||
permission_manager=self.permission_manager,
|
||||
@@ -159,6 +229,7 @@ class HanzoMCPServer:
|
||||
enabled_tools=final_enabled_tools,
|
||||
)
|
||||
|
||||
register_all_prompts = _get_register_all_prompts()
|
||||
register_all_prompts(mcp_server=self.mcp, projects=self.project_paths)
|
||||
|
||||
def _setup_cleanup_handlers(self) -> None:
|
||||
@@ -192,6 +263,7 @@ class HanzoMCPServer:
|
||||
|
||||
def _background_cleanup(self) -> None:
|
||||
"""Background thread for periodic session cleanup."""
|
||||
SessionStorage = _get_session_storage()
|
||||
while not self._shutdown_event.is_set():
|
||||
try:
|
||||
# Clean up expired sessions every 2 minutes
|
||||
@@ -207,6 +279,7 @@ class HanzoMCPServer:
|
||||
def _cleanup_sessions(self) -> None:
|
||||
"""Clean up all active sessions."""
|
||||
try:
|
||||
SessionStorage = _get_session_storage()
|
||||
cleared_count = SessionStorage.clear_all_sessions()
|
||||
if cleared_count > 0:
|
||||
# Only log if not stdio transport
|
||||
|
||||
@@ -5,21 +5,123 @@ 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
|
||||
|
||||
IMPORTANT: All heavy imports are LAZY to ensure fast MCP startup (<1 second).
|
||||
Imports only happen when register_all_tools() is called, not at module load time.
|
||||
"""
|
||||
|
||||
# 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
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
class FastMCP: # type: ignore
|
||||
pass
|
||||
# TYPE_CHECKING imports - these don't execute at runtime
|
||||
if TYPE_CHECKING:
|
||||
from mcp.server import FastMCP
|
||||
from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.permissions import PermissionManager
|
||||
|
||||
# Flags for optional tool availability - checked lazily
|
||||
_LSP_TOOL_AVAILABLE: bool | None = None
|
||||
_REFACTOR_TOOL_AVAILABLE: bool | None = None
|
||||
_MEMORY_TOOLS_AVAILABLE: bool | None = None
|
||||
|
||||
|
||||
try: # pragma: no cover
|
||||
def _check_lsp_available() -> bool:
|
||||
"""Check if LSP tool is available (lazy check)."""
|
||||
global _LSP_TOOL_AVAILABLE
|
||||
if _LSP_TOOL_AVAILABLE is None:
|
||||
try:
|
||||
from hanzo_mcp.tools.lsp import LSPTool # noqa: F401
|
||||
_LSP_TOOL_AVAILABLE = True
|
||||
except ImportError:
|
||||
_LSP_TOOL_AVAILABLE = False
|
||||
return _LSP_TOOL_AVAILABLE
|
||||
|
||||
|
||||
def _check_refactor_available() -> bool:
|
||||
"""Check if refactor tool is available (lazy check)."""
|
||||
global _REFACTOR_TOOL_AVAILABLE
|
||||
if _REFACTOR_TOOL_AVAILABLE is None:
|
||||
try:
|
||||
from hanzo_mcp.tools.refactor import RefactorTool # noqa: F401
|
||||
_REFACTOR_TOOL_AVAILABLE = True
|
||||
except ImportError:
|
||||
_REFACTOR_TOOL_AVAILABLE = False
|
||||
return _REFACTOR_TOOL_AVAILABLE
|
||||
|
||||
|
||||
def _check_memory_available() -> bool:
|
||||
"""Check if memory tools are available (lazy check)."""
|
||||
global _MEMORY_TOOLS_AVAILABLE
|
||||
if _MEMORY_TOOLS_AVAILABLE is None:
|
||||
try:
|
||||
from hanzo_mcp.tools.memory import register_memory_tools # noqa: F401
|
||||
_MEMORY_TOOLS_AVAILABLE = True
|
||||
except ImportError:
|
||||
_MEMORY_TOOLS_AVAILABLE = False
|
||||
return _MEMORY_TOOLS_AVAILABLE
|
||||
|
||||
|
||||
# Expose availability flags as properties for backward compatibility
|
||||
@property
|
||||
def LSP_TOOL_AVAILABLE() -> bool:
|
||||
return _check_lsp_available()
|
||||
|
||||
|
||||
@property
|
||||
def REFACTOR_TOOL_AVAILABLE() -> bool:
|
||||
return _check_refactor_available()
|
||||
|
||||
|
||||
@property
|
||||
def MEMORY_TOOLS_AVAILABLE() -> bool:
|
||||
return _check_memory_available()
|
||||
|
||||
|
||||
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
|
||||
"""
|
||||
# LAZY IMPORTS - Only import when this function is called, not at module load time
|
||||
# This is critical for fast MCP startup (<1 second)
|
||||
import logging
|
||||
|
||||
from hanzo_mcp.tools.llm import (
|
||||
LLMTool,
|
||||
ConsensusTool,
|
||||
@@ -51,109 +153,21 @@ try: # pragma: no cover
|
||||
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
|
||||
from hanzo_mcp.tools.common.plugin_loader import load_user_plugins
|
||||
|
||||
# Try memory tools
|
||||
try:
|
||||
from hanzo_mcp.tools.memory import register_memory_tools
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
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
|
||||
# Define stub functions for required imports
|
||||
def activate_mode_from_env():
|
||||
pass
|
||||
|
||||
class ModeLoader:
|
||||
@staticmethod
|
||||
def get_enabled_tools_from_mode(base_enabled_tools=None, force_mode=None):
|
||||
return base_enabled_tools or {}
|
||||
|
||||
@staticmethod
|
||||
def apply_environment_from_mode():
|
||||
pass
|
||||
|
||||
def load_user_plugins():
|
||||
return {}
|
||||
|
||||
|
||||
# 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
|
||||
|
||||
# Try to import refactor tool
|
||||
try:
|
||||
from hanzo_mcp.tools.refactor import RefactorTool, create_refactor_tool
|
||||
|
||||
REFACTOR_TOOL_AVAILABLE = True
|
||||
except ImportError:
|
||||
REFACTOR_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 = {}
|
||||
|
||||
@@ -212,9 +226,6 @@ def register_all_tools(
|
||||
)
|
||||
# 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(
|
||||
@@ -466,8 +477,9 @@ def register_all_tools(
|
||||
"manage_knowledge_bases": is_tool_enabled("manage_knowledge_bases", True),
|
||||
}
|
||||
|
||||
if any(memory_enabled.values()) and MEMORY_TOOLS_AVAILABLE:
|
||||
if any(memory_enabled.values()) and _check_memory_available():
|
||||
try:
|
||||
from hanzo_mcp.tools.memory import register_memory_tools
|
||||
memory_tools = register_memory_tools(
|
||||
mcp_server, permission_manager, user_id="default", project_id="default"
|
||||
)
|
||||
@@ -479,8 +491,9 @@ def register_all_tools(
|
||||
logger.warning(f"Failed to register memory tools: {e}")
|
||||
|
||||
# Register LSP tool if enabled
|
||||
if is_tool_enabled("lsp", True) and LSP_TOOL_AVAILABLE:
|
||||
if is_tool_enabled("lsp", True) and _check_lsp_available():
|
||||
try:
|
||||
from hanzo_mcp.tools.lsp import create_lsp_tool
|
||||
tool = create_lsp_tool()
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
@@ -488,8 +501,9 @@ def register_all_tools(
|
||||
logger.warning(f"Failed to register LSP tool: {e}")
|
||||
|
||||
# Register refactor tool if enabled
|
||||
if is_tool_enabled("refactor", True) and REFACTOR_TOOL_AVAILABLE:
|
||||
if is_tool_enabled("refactor", True) and _check_refactor_available():
|
||||
try:
|
||||
from hanzo_mcp.tools.refactor import create_refactor_tool
|
||||
tool = create_refactor_tool()
|
||||
tool.register(mcp_server)
|
||||
all_tools[tool.name] = tool
|
||||
|
||||
@@ -2,9 +2,12 @@
|
||||
|
||||
These tools use the hanzo-memory package to manage knowledge bases and facts,
|
||||
supporting hierarchical organization (session, project, global).
|
||||
|
||||
IMPORTANT: All hanzo-memory imports are lazy to avoid slow startup.
|
||||
The embedding service initialization takes 3+ seconds which would block MCP startup.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, final, override
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, final, override
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
@@ -13,18 +16,12 @@ from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.auto_timeout import auto_timeout
|
||||
|
||||
# Import from hanzo-memory package
|
||||
try:
|
||||
from hanzo_memory.models.memory import Memory, MemoryWithScore
|
||||
from hanzo_memory.services.memory import MemoryService, get_memory_service
|
||||
from hanzo_memory.models.knowledge import Fact, FactCreate, KnowledgeBase
|
||||
# Type hints only - no runtime import
|
||||
if TYPE_CHECKING:
|
||||
from hanzo_memory.services.memory import MemoryService
|
||||
|
||||
KNOWLEDGE_AVAILABLE = True
|
||||
except ImportError:
|
||||
KNOWLEDGE_AVAILABLE = False
|
||||
raise ImportError(
|
||||
"hanzo-memory package is required for knowledge tools. Install it from ~/work/hanzo/ide/pkg/memory"
|
||||
)
|
||||
# Use lazy loading from memory_tools
|
||||
from hanzo_mcp.tools.memory.memory_tools import _check_memory_available, _get_lazy_memory_service
|
||||
|
||||
|
||||
class KnowledgeToolBase(BaseTool):
|
||||
@@ -40,8 +37,15 @@ class KnowledgeToolBase(BaseTool):
|
||||
"""
|
||||
self.user_id = user_id
|
||||
self.project_id = project_id
|
||||
# Use the memory service for knowledge management
|
||||
self.service = get_memory_service()
|
||||
# Lazy service loading - don't initialize until first use
|
||||
self._service: Optional["MemoryService"] = None
|
||||
|
||||
@property
|
||||
def service(self) -> "MemoryService":
|
||||
"""Get memory service lazily."""
|
||||
if self._service is None:
|
||||
self._service = _get_lazy_memory_service()
|
||||
return self._service
|
||||
|
||||
|
||||
@final
|
||||
@@ -319,10 +323,8 @@ summarize_to_memory(content="Company guidelines...", topic="Guidelines", scope="
|
||||
# This would typically use an LLM to summarize, but for now we'll store as-is
|
||||
summary = f"Summary of {topic}:\n{content[:500]}..." if len(content) > 500 else content
|
||||
|
||||
# Store the summary as a memory
|
||||
from hanzo_memory.services.memory import get_memory_service
|
||||
|
||||
memory_service = get_memory_service()
|
||||
# Store the summary as a memory (using lazy service)
|
||||
memory_service = self.service
|
||||
|
||||
# Determine scope
|
||||
if scope == "global":
|
||||
|
||||
@@ -2,9 +2,12 @@
|
||||
|
||||
This module provides MCP tools that use the hanzo-memory package as a library.
|
||||
The hanzo-memory package provides embedded database and vector search capabilities.
|
||||
|
||||
IMPORTANT: All hanzo-memory imports are lazy to avoid slow startup.
|
||||
The embedding service initialization takes 3+ seconds which would block MCP startup.
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional, final, override
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, final, override
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
@@ -13,17 +16,40 @@ from hanzo_mcp.tools.common.base import BaseTool
|
||||
from hanzo_mcp.tools.common.context import create_tool_context
|
||||
from hanzo_mcp.tools.common.auto_timeout import auto_timeout
|
||||
|
||||
# Import from hanzo-memory package
|
||||
try:
|
||||
# Type hints only - no runtime import
|
||||
if TYPE_CHECKING:
|
||||
from hanzo_memory.models.memory import Memory, MemoryWithScore
|
||||
from hanzo_memory.services.memory import MemoryService, get_memory_service
|
||||
from hanzo_memory.services.memory import MemoryService
|
||||
|
||||
MEMORY_AVAILABLE = True
|
||||
except ImportError:
|
||||
MEMORY_AVAILABLE = False
|
||||
raise ImportError(
|
||||
"hanzo-memory package is required for memory tools. Install it from ~/work/hanzo/python-sdk/pkg/hanzo-memory"
|
||||
)
|
||||
# Lazy loading - don't import at module load time
|
||||
MEMORY_AVAILABLE: Optional[bool] = None
|
||||
_memory_service: Optional["MemoryService"] = None
|
||||
|
||||
|
||||
def _check_memory_available() -> bool:
|
||||
"""Check if hanzo-memory is available (lazy check)."""
|
||||
global MEMORY_AVAILABLE
|
||||
if MEMORY_AVAILABLE is None:
|
||||
try:
|
||||
import hanzo_memory # noqa: F401
|
||||
MEMORY_AVAILABLE = True
|
||||
except ImportError:
|
||||
MEMORY_AVAILABLE = False
|
||||
return MEMORY_AVAILABLE
|
||||
|
||||
|
||||
def _get_lazy_memory_service() -> "MemoryService":
|
||||
"""Get memory service lazily - only import when actually needed."""
|
||||
global _memory_service
|
||||
if _memory_service is None:
|
||||
if not _check_memory_available():
|
||||
raise ImportError(
|
||||
"hanzo-memory package is required for memory tools. "
|
||||
"Install with: pip install hanzo-memory"
|
||||
)
|
||||
from hanzo_memory.services.memory import get_memory_service
|
||||
_memory_service = get_memory_service()
|
||||
return _memory_service
|
||||
|
||||
|
||||
class MemoryToolBase(BaseTool):
|
||||
@@ -39,8 +65,15 @@ class MemoryToolBase(BaseTool):
|
||||
"""
|
||||
self.user_id = user_id
|
||||
self.project_id = project_id
|
||||
# Use the global memory service from hanzo-memory
|
||||
self.service = get_memory_service()
|
||||
# Lazy service loading - don't initialize until first use
|
||||
self._service: Optional["MemoryService"] = None
|
||||
|
||||
@property
|
||||
def service(self) -> "MemoryService":
|
||||
"""Get memory service lazily."""
|
||||
if self._service is None:
|
||||
self._service = _get_lazy_memory_service()
|
||||
return self._service
|
||||
|
||||
|
||||
@final
|
||||
|
||||
@@ -6,7 +6,6 @@ search capabilities including text, AST, symbols, memory, and semantic search.
|
||||
|
||||
import json
|
||||
import time
|
||||
import hashlib
|
||||
import subprocess
|
||||
from typing import Any, Dict, List, Optional
|
||||
from pathlib import Path
|
||||
@@ -31,13 +30,8 @@ try:
|
||||
except ImportError:
|
||||
TREESITTER_AVAILABLE = False
|
||||
|
||||
try:
|
||||
import chromadb
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
VECTOR_SEARCH_AVAILABLE = True
|
||||
except ImportError:
|
||||
VECTOR_SEARCH_AVAILABLE = False
|
||||
# Vector search removed - too heavy for MCP tools
|
||||
# If semantic search is needed, use external services (hanzo-node, hanzo desktop)
|
||||
|
||||
# LSP availability check
|
||||
try:
|
||||
@@ -47,6 +41,14 @@ try:
|
||||
except ImportError:
|
||||
LSP_AVAILABLE = False
|
||||
|
||||
# Git search availability
|
||||
try:
|
||||
from hanzo_mcp.tools.filesystem.git_search import GitSearchTool
|
||||
|
||||
GIT_SEARCH_AVAILABLE = True
|
||||
except ImportError:
|
||||
GIT_SEARCH_AVAILABLE = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchResult:
|
||||
@@ -93,7 +95,6 @@ class SearchTool(BaseTool):
|
||||
- Symbol definitions and references (using ctags/LSP)
|
||||
- Files and directories (using find tool)
|
||||
- Memory and knowledge base entries
|
||||
- Semantic/conceptual matches (using vector search)
|
||||
|
||||
The tool automatically determines the best search strategy based on your query
|
||||
and runs multiple search types in parallel for comprehensive results.
|
||||
@@ -114,11 +115,7 @@ class SearchTool(BaseTool):
|
||||
search("test_*.py", search_files=True) # Find test files
|
||||
search("config", search_files=True) # Find config files
|
||||
|
||||
4. Semantic search:
|
||||
search("how authentication works") # Natural language query
|
||||
search("database connection logic") # Conceptual search
|
||||
|
||||
5. Memory search:
|
||||
4. Memory search:
|
||||
search("previous discussion about API design") # Search memories
|
||||
search("that bug we fixed last week") # Search knowledge
|
||||
|
||||
@@ -131,7 +128,6 @@ class SearchTool(BaseTool):
|
||||
- Respects .gitignore and other exclusions
|
||||
|
||||
PRO TIPS:
|
||||
- Use natural language for conceptual searches
|
||||
- Use code syntax for exact matches
|
||||
- Add search_files=True to also find filenames
|
||||
- Results are ranked by relevance and type
|
||||
@@ -139,26 +135,17 @@ class SearchTool(BaseTool):
|
||||
"""
|
||||
|
||||
name = "search"
|
||||
description = """THE primary unified search tool for rapid parallel search across all modalities.
|
||||
|
||||
Find anything in your codebase using text, AST, symbols, files, memory, and semantic search.
|
||||
Automatically detects query intent and runs appropriate searches in parallel.
|
||||
description = """THE primary unified search tool for rapid parallel search.
|
||||
|
||||
Find anything in your codebase using text, AST, symbols, files, and memory search.
|
||||
Fast and lightweight - no heavy ML dependencies.
|
||||
"""
|
||||
|
||||
def __init__(self, enable_vector_index: bool = False):
|
||||
"""Initialize search tool.
|
||||
|
||||
Args:
|
||||
enable_vector_index: Whether to enable vector search indexing (default: False).
|
||||
When False, vector search is disabled for faster startup.
|
||||
Set to True to enable semantic search capabilities.
|
||||
"""
|
||||
def __init__(self):
|
||||
"""Initialize search tool - fast and lightweight."""
|
||||
super().__init__()
|
||||
self.ripgrep_available = self._check_ripgrep()
|
||||
self.vector_db = None
|
||||
self.embedder = None
|
||||
self._vector_initialized = False
|
||||
self._enable_vector_index = enable_vector_index
|
||||
self.git_available = self._check_git()
|
||||
|
||||
def _check_ripgrep(self) -> bool:
|
||||
"""Check if ripgrep is available."""
|
||||
@@ -168,47 +155,14 @@ class SearchTool(BaseTool):
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def _init_vector_search(self) -> bool:
|
||||
"""Initialize vector search components lazily.
|
||||
|
||||
Returns:
|
||||
True if initialization succeeded, False otherwise.
|
||||
"""
|
||||
if self._vector_initialized:
|
||||
return self.vector_db is not None
|
||||
|
||||
self._vector_initialized = True
|
||||
|
||||
if not VECTOR_SEARCH_AVAILABLE:
|
||||
return False
|
||||
|
||||
if not self._enable_vector_index:
|
||||
return False
|
||||
|
||||
def _check_git(self) -> bool:
|
||||
"""Check if git is available."""
|
||||
try:
|
||||
self.embedder = SentenceTransformer("all-MiniLM-L6-v2")
|
||||
self.vector_db = chromadb.Client()
|
||||
# Create or get collection
|
||||
self.collection = self.vector_db.get_or_create_collection(
|
||||
name="code_search", metadata={"description": "Code semantic search"}
|
||||
)
|
||||
subprocess.run(["git", "--version"], capture_output=True, check=True)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Failed to initialize vector search: {e}")
|
||||
self.vector_db = None
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def _should_use_vector_search(self, query: str) -> bool:
|
||||
"""Determine if vector search would be helpful."""
|
||||
# Use vector search for natural language queries
|
||||
indicators = [
|
||||
len(query.split()) > 2, # Multi-word queries
|
||||
not any(c in query for c in ["(", ")", "{", "}", "[", "]"]), # Not code syntax
|
||||
" " in query, # Has spaces (natural language)
|
||||
not query.startswith("^") and not query.endswith("$"), # Not regex anchors
|
||||
]
|
||||
return sum(indicators) >= 2
|
||||
|
||||
def _should_use_ast_search(self, query: str) -> bool:
|
||||
"""Determine if AST search would be helpful."""
|
||||
# Use AST search for code patterns
|
||||
@@ -229,6 +183,19 @@ class SearchTool(BaseTool):
|
||||
and not " " in query.strip() # Single token
|
||||
)
|
||||
|
||||
def _is_git_repo(self, path: str) -> bool:
|
||||
"""Check if path is inside a git repository."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", "rev-parse", "--git-dir"],
|
||||
capture_output=True,
|
||||
cwd=path or ".",
|
||||
timeout=2,
|
||||
)
|
||||
return result.returncode == 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def run(
|
||||
self,
|
||||
pattern: str,
|
||||
@@ -239,9 +206,9 @@ class SearchTool(BaseTool):
|
||||
context_lines: int = 3,
|
||||
search_files: bool = False,
|
||||
search_memory: bool = None,
|
||||
search_git_history: bool = None,
|
||||
enable_text: bool = None,
|
||||
enable_ast: bool = None,
|
||||
enable_vector: bool = None,
|
||||
enable_symbol: bool = None,
|
||||
enable_lsp: bool = None,
|
||||
page_size: int = 50,
|
||||
@@ -259,6 +226,7 @@ class SearchTool(BaseTool):
|
||||
context_lines: Lines of context around text matches
|
||||
search_files: Also search for matching filenames
|
||||
search_memory: Search in memory/knowledge base (auto-detected if None)
|
||||
search_git_history: Search git commit history (auto-detected if None)
|
||||
enable_*: Force enable/disable specific search types (auto if None)
|
||||
enable_lsp: Enable LSP-based reference search (auto if None)
|
||||
page_size: Results per page (default: 50)
|
||||
@@ -268,22 +236,13 @@ class SearchTool(BaseTool):
|
||||
# Auto-detect search types based on query
|
||||
if search_memory is None:
|
||||
# Search memory for natural language queries or specific references
|
||||
search_memory = MEMORY_AVAILABLE and (
|
||||
self._should_use_vector_search(pattern)
|
||||
or any(word in pattern.lower() for word in ["previous", "discussion", "remember", "last"])
|
||||
search_memory = MEMORY_AVAILABLE and any(
|
||||
word in pattern.lower() for word in ["previous", "discussion", "remember", "last"]
|
||||
)
|
||||
|
||||
if enable_text is None:
|
||||
enable_text = True # Always use text search as baseline
|
||||
|
||||
if enable_vector is None:
|
||||
# Only enable vector search if explicitly enabled via enable_vector_index
|
||||
enable_vector = (
|
||||
self._enable_vector_index
|
||||
and self._should_use_vector_search(pattern)
|
||||
and VECTOR_SEARCH_AVAILABLE
|
||||
)
|
||||
|
||||
if enable_ast is None:
|
||||
enable_ast = self._should_use_ast_search(pattern) and TREESITTER_AVAILABLE
|
||||
|
||||
@@ -294,6 +253,10 @@ class SearchTool(BaseTool):
|
||||
# Use LSP for symbol-like queries when LSP is available
|
||||
enable_lsp = LSP_AVAILABLE and self._should_use_symbol_search(pattern)
|
||||
|
||||
if search_git_history is None:
|
||||
# Auto-enable git history for identifier-like patterns in git repos
|
||||
search_git_history = (GIT_SEARCH_AVAILABLE or self.git_available) and self._is_git_repo(path) and self._should_use_symbol_search(pattern)
|
||||
|
||||
# Collect results from all enabled search types
|
||||
all_results = []
|
||||
search_stats = {
|
||||
@@ -305,63 +268,76 @@ class SearchTool(BaseTool):
|
||||
"time_ms": {},
|
||||
}
|
||||
|
||||
# 1. Text search (ripgrep) - always fast, do first
|
||||
# Run all searches in PARALLEL using asyncio.gather
|
||||
import asyncio
|
||||
|
||||
async def timed_search(name: str, coro):
|
||||
"""Wrapper to time each search."""
|
||||
start = time.time()
|
||||
try:
|
||||
results = await coro
|
||||
except Exception as e:
|
||||
print(f"{name} search error: {e}")
|
||||
results = []
|
||||
elapsed = int((time.time() - start) * 1000)
|
||||
return name, results, elapsed
|
||||
|
||||
# Build list of search coroutines to run in parallel
|
||||
search_tasks = []
|
||||
|
||||
if enable_text:
|
||||
start = time.time()
|
||||
text_results = await self._text_search(pattern, path, include, exclude, max_results_per_type, context_lines)
|
||||
search_stats["time_ms"]["text"] = int((time.time() - start) * 1000)
|
||||
search_stats["search_types_used"].append("text")
|
||||
all_results.extend(text_results)
|
||||
|
||||
# 2. AST search - for code structure
|
||||
search_tasks.append(timed_search(
|
||||
"text",
|
||||
self._text_search(pattern, path, include, exclude, max_results_per_type, context_lines)
|
||||
))
|
||||
|
||||
if enable_ast and TREESITTER_AVAILABLE:
|
||||
start = time.time()
|
||||
ast_results = await self._ast_search(pattern, path, include, exclude, max_results_per_type, context_lines)
|
||||
search_stats["time_ms"]["ast"] = int((time.time() - start) * 1000)
|
||||
search_stats["search_types_used"].append("ast")
|
||||
all_results.extend(ast_results)
|
||||
|
||||
# 3. Symbol search - for definitions
|
||||
search_tasks.append(timed_search(
|
||||
"ast",
|
||||
self._ast_search(pattern, path, include, exclude, max_results_per_type, context_lines)
|
||||
))
|
||||
|
||||
if enable_symbol:
|
||||
start = time.time()
|
||||
symbol_results = await self._symbol_search(pattern, path, include, exclude, max_results_per_type)
|
||||
search_stats["time_ms"]["symbol"] = int((time.time() - start) * 1000)
|
||||
search_stats["search_types_used"].append("symbol")
|
||||
all_results.extend(symbol_results)
|
||||
|
||||
# 4. Vector search - for semantic similarity (lazy init)
|
||||
if enable_vector and self._init_vector_search():
|
||||
start = time.time()
|
||||
vector_results = await self._vector_search(
|
||||
pattern, path, include, exclude, max_results_per_type, context_lines
|
||||
)
|
||||
search_stats["time_ms"]["vector"] = int((time.time() - start) * 1000)
|
||||
search_stats["search_types_used"].append("vector")
|
||||
all_results.extend(vector_results)
|
||||
|
||||
# 5. File search - for finding files by name/pattern
|
||||
search_tasks.append(timed_search(
|
||||
"symbol",
|
||||
self._symbol_search(pattern, path, include, exclude, max_results_per_type)
|
||||
))
|
||||
|
||||
if search_files:
|
||||
start = time.time()
|
||||
file_results = await self._file_search(pattern, path, include, exclude, max_results_per_type)
|
||||
search_stats["time_ms"]["files"] = int((time.time() - start) * 1000)
|
||||
search_stats["search_types_used"].append("files")
|
||||
all_results.extend(file_results)
|
||||
|
||||
# 6. Memory search - for knowledge base and previous discussions
|
||||
search_tasks.append(timed_search(
|
||||
"files",
|
||||
self._file_search(pattern, path, include, exclude, max_results_per_type)
|
||||
))
|
||||
|
||||
if search_memory:
|
||||
start = time.time()
|
||||
memory_results = await self._memory_search(pattern, max_results_per_type, context_lines)
|
||||
search_stats["time_ms"]["memory"] = int((time.time() - start) * 1000)
|
||||
search_stats["search_types_used"].append("memory")
|
||||
all_results.extend(memory_results)
|
||||
|
||||
# 7. LSP search - for precise symbol references and definitions
|
||||
search_tasks.append(timed_search(
|
||||
"memory",
|
||||
self._memory_search(pattern, max_results_per_type, context_lines)
|
||||
))
|
||||
|
||||
if enable_lsp and LSP_AVAILABLE:
|
||||
start = time.time()
|
||||
lsp_results = await self._lsp_search(pattern, path, include, max_results_per_type)
|
||||
search_stats["time_ms"]["lsp"] = int((time.time() - start) * 1000)
|
||||
search_stats["search_types_used"].append("lsp")
|
||||
all_results.extend(lsp_results)
|
||||
search_tasks.append(timed_search(
|
||||
"lsp",
|
||||
self._lsp_search(pattern, path, include, max_results_per_type)
|
||||
))
|
||||
|
||||
if search_git_history and self.git_available:
|
||||
search_tasks.append(timed_search(
|
||||
"git",
|
||||
self._git_history_search(pattern, path, max_results_per_type)
|
||||
))
|
||||
|
||||
# Execute all searches in parallel
|
||||
if search_tasks:
|
||||
results_list = await asyncio.gather(*search_tasks, return_exceptions=True)
|
||||
|
||||
for result in results_list:
|
||||
if isinstance(result, Exception):
|
||||
continue
|
||||
name, results, elapsed = result
|
||||
search_stats["time_ms"][name] = elapsed
|
||||
search_stats["search_types_used"].append(name)
|
||||
all_results.extend(results)
|
||||
|
||||
# Deduplicate and rank results
|
||||
unique_results = self._deduplicate_results(all_results)
|
||||
@@ -421,10 +397,10 @@ class SearchTool(BaseTool):
|
||||
page: int = 1,
|
||||
enable_text: bool = True,
|
||||
enable_ast: bool = True,
|
||||
enable_vector: bool = True,
|
||||
enable_symbol: bool = True,
|
||||
search_files: bool = False,
|
||||
search_memory: bool = False,
|
||||
search_git_history: bool = False,
|
||||
) -> str:
|
||||
"""Execute unified search."""
|
||||
return await self.call(
|
||||
@@ -438,10 +414,10 @@ class SearchTool(BaseTool):
|
||||
page=page,
|
||||
enable_text=enable_text,
|
||||
enable_ast=enable_ast,
|
||||
enable_vector=enable_vector,
|
||||
enable_symbol=enable_symbol,
|
||||
search_files=search_files,
|
||||
search_memory=search_memory,
|
||||
search_git_history=search_git_history,
|
||||
)
|
||||
|
||||
async def _text_search(
|
||||
@@ -530,7 +506,10 @@ class SearchTool(BaseTool):
|
||||
max_results: int,
|
||||
context_lines: int,
|
||||
) -> List[SearchResult]:
|
||||
"""Perform AST-based search using treesitter."""
|
||||
"""Perform AST-based search using treesitter.
|
||||
|
||||
OPTIMIZED: Uses ripgrep first to narrow down files before tree-sitter parsing.
|
||||
"""
|
||||
# Try to use grep-ast if available
|
||||
try:
|
||||
from grep_ast.grep_ast import TreeContext
|
||||
@@ -541,23 +520,49 @@ class SearchTool(BaseTool):
|
||||
results = []
|
||||
|
||||
try:
|
||||
# Get files to search
|
||||
search_path = Path(path or ".")
|
||||
files_to_search = []
|
||||
|
||||
if search_path.is_file():
|
||||
files_to_search = [search_path]
|
||||
else:
|
||||
# Find files matching include pattern
|
||||
pattern_to_use = include or "*.py"
|
||||
for ext in ["*.py", "*.js", "*.ts", "*.go", "*.java", "*.cpp", "*.c"]:
|
||||
if include and include != ext:
|
||||
continue
|
||||
files_to_search.extend(search_path.rglob(ext))
|
||||
if len(files_to_search) >= max_results:
|
||||
break
|
||||
# OPTIMIZATION: Use ripgrep first to find files containing the pattern
|
||||
# This is MUCH faster than scanning all files with rglob + tree-sitter
|
||||
if self.ripgrep_available:
|
||||
cmd = ["rg", "--files-with-matches", "-l", "--max-count", "1"]
|
||||
|
||||
# Add file type filters
|
||||
if include:
|
||||
cmd.extend(["--glob", include])
|
||||
else:
|
||||
# Default to common source files
|
||||
for ext in ["*.py", "*.js", "*.ts", "*.go", "*.java", "*.cpp", "*.c", "*.rs"]:
|
||||
cmd.extend(["--glob", ext])
|
||||
|
||||
if exclude:
|
||||
cmd.extend(["--glob", f"!{exclude}"])
|
||||
|
||||
cmd.extend([pattern, str(search_path)])
|
||||
|
||||
try:
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=5)
|
||||
if result.returncode == 0 and result.stdout.strip():
|
||||
files_to_search = [Path(f) for f in result.stdout.strip().split("\n") if f]
|
||||
except (subprocess.TimeoutExpired, subprocess.SubprocessError):
|
||||
pass
|
||||
|
||||
# Fallback to limited rglob if ripgrep didn't find anything or isn't available
|
||||
if not files_to_search:
|
||||
exts = [include] if include else ["*.py", "*.js", "*.ts", "*.go"]
|
||||
for ext in exts[:3]: # Limit extensions to check
|
||||
for f in search_path.rglob(ext):
|
||||
files_to_search.append(f)
|
||||
if len(files_to_search) >= max_results:
|
||||
break
|
||||
if len(files_to_search) >= max_results:
|
||||
break
|
||||
|
||||
# Search each file
|
||||
# Search each file with tree-sitter (now limited to files that actually contain pattern)
|
||||
for file_path in files_to_search[:max_results]:
|
||||
if not file_path.is_file():
|
||||
continue
|
||||
@@ -644,54 +649,6 @@ class SearchTool(BaseTool):
|
||||
|
||||
return results
|
||||
|
||||
async def _vector_search(
|
||||
self,
|
||||
query: str,
|
||||
path: str,
|
||||
include: Optional[str],
|
||||
exclude: Optional[str],
|
||||
max_results: int,
|
||||
context_lines: int,
|
||||
) -> List[SearchResult]:
|
||||
"""Perform semantic vector search."""
|
||||
if not self.vector_db or not self.embedder:
|
||||
return []
|
||||
|
||||
results = []
|
||||
|
||||
try:
|
||||
# Embed the query
|
||||
query_embedding = self.embedder.encode(query).tolist()
|
||||
|
||||
# Search in vector database
|
||||
search_results = self.collection.query(
|
||||
query_embeddings=[query_embedding],
|
||||
n_results=max_results,
|
||||
where={"path": {"$contains": path}} if path != "." else None,
|
||||
)
|
||||
|
||||
if search_results["ids"][0]:
|
||||
for i, _doc_id in enumerate(search_results["ids"][0]):
|
||||
metadata = search_results["metadatas"][0][i]
|
||||
|
||||
result = SearchResult(
|
||||
file_path=metadata["file_path"],
|
||||
line_number=metadata["line_number"],
|
||||
column=0,
|
||||
match_text=search_results["documents"][0][i],
|
||||
context_before=[],
|
||||
context_after=[],
|
||||
match_type="vector",
|
||||
score=1.0 - search_results["distances"][0][i], # Convert distance to similarity
|
||||
semantic_context=metadata.get("context", ""),
|
||||
)
|
||||
results.append(result)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Vector search error: {e}")
|
||||
|
||||
return results
|
||||
|
||||
async def _file_search(
|
||||
self,
|
||||
pattern: str,
|
||||
@@ -908,6 +865,102 @@ class SearchTool(BaseTool):
|
||||
|
||||
return results
|
||||
|
||||
async def _git_history_search(
|
||||
self,
|
||||
pattern: str,
|
||||
path: str,
|
||||
max_results: int,
|
||||
) -> List[SearchResult]:
|
||||
"""Search git commit history for pattern using git log -S (pickaxe)."""
|
||||
results = []
|
||||
|
||||
if not self.git_available:
|
||||
return results
|
||||
|
||||
try:
|
||||
# Use git log -S to find commits that added/removed the pattern
|
||||
# This is much faster than searching through all commits
|
||||
cmd = [
|
||||
"git", "log",
|
||||
f"-S{pattern}", # Search for commits that add/remove pattern
|
||||
"--oneline",
|
||||
f"-n{max_results}",
|
||||
"--pretty=format:%h|%s|%an|%ar",
|
||||
]
|
||||
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=path or ".",
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
if result.returncode == 0 and result.stdout.strip():
|
||||
for line in result.stdout.strip().split("\n"):
|
||||
parts = line.split("|", 3)
|
||||
if len(parts) >= 4:
|
||||
commit_hash, subject, author, date = parts
|
||||
|
||||
# Create a result for this commit
|
||||
search_result = SearchResult(
|
||||
file_path=f"git://{commit_hash}",
|
||||
line_number=0,
|
||||
column=0,
|
||||
match_text=subject.strip(),
|
||||
context_before=[],
|
||||
context_after=[],
|
||||
match_type="git",
|
||||
score=0.85,
|
||||
semantic_context=f"Commit by {author} ({date})",
|
||||
)
|
||||
results.append(search_result)
|
||||
|
||||
# Also try git grep for current HEAD if we don't have many results
|
||||
if len(results) < max_results // 2:
|
||||
grep_cmd = [
|
||||
"git", "grep",
|
||||
"-n", # Line numbers
|
||||
"-I", # Skip binary files
|
||||
f"--max-count={max_results - len(results)}",
|
||||
pattern,
|
||||
]
|
||||
|
||||
grep_result = subprocess.run(
|
||||
grep_cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=path or ".",
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
if grep_result.returncode == 0 and grep_result.stdout.strip():
|
||||
for line in grep_result.stdout.strip().split("\n")[:max_results - len(results)]:
|
||||
# Format: file:line:content
|
||||
parts = line.split(":", 2)
|
||||
if len(parts) >= 3:
|
||||
file_path, line_num, content = parts
|
||||
|
||||
search_result = SearchResult(
|
||||
file_path=str(Path(path or ".").resolve() / file_path),
|
||||
line_number=int(line_num) if line_num.isdigit() else 1,
|
||||
column=0,
|
||||
match_text=content.strip(),
|
||||
context_before=[],
|
||||
context_after=[],
|
||||
match_type="git",
|
||||
score=0.9,
|
||||
semantic_context="git grep match",
|
||||
)
|
||||
results.append(search_result)
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
print("Git history search timed out")
|
||||
except Exception as e:
|
||||
print(f"Git history search error: {e}")
|
||||
|
||||
return results
|
||||
|
||||
def _deduplicate_results(self, results: List[SearchResult]) -> List[SearchResult]:
|
||||
"""Remove duplicate results across search types."""
|
||||
seen = set()
|
||||
@@ -1042,98 +1095,7 @@ class SearchTool(BaseTool):
|
||||
return results
|
||||
|
||||
|
||||
# Index builder for vector search
|
||||
class CodeIndexer:
|
||||
"""Build and maintain vector search index."""
|
||||
|
||||
def __init__(self, vector_db, embedder):
|
||||
self.vector_db = vector_db
|
||||
self.embedder = embedder
|
||||
self.collection = vector_db.get_or_create_collection("code_search")
|
||||
|
||||
async def index_directory(self, path: str, file_patterns: List[str] = None):
|
||||
"""Index a directory for vector search."""
|
||||
if file_patterns is None:
|
||||
file_patterns = ["*.py", "*.js", "*.ts", "*.go", "*.java", "*.cpp", "*.c"]
|
||||
|
||||
documents = []
|
||||
metadatas = []
|
||||
ids = []
|
||||
|
||||
for pattern in file_patterns:
|
||||
for file_path in Path(path).rglob(pattern):
|
||||
if file_path.is_file():
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
|
||||
# Split into chunks (functions, classes, etc.)
|
||||
chunks = self._split_code_intelligently(content, file_path)
|
||||
|
||||
for chunk in chunks:
|
||||
doc_id = hashlib.md5(
|
||||
f"{file_path}:{chunk['line']}:{chunk['text'][:50]}".encode()
|
||||
).hexdigest()
|
||||
|
||||
documents.append(chunk["text"])
|
||||
metadatas.append(
|
||||
{
|
||||
"file_path": str(file_path),
|
||||
"line_number": chunk["line"],
|
||||
"context": chunk.get("context", ""),
|
||||
"type": chunk.get("type", "code"),
|
||||
}
|
||||
)
|
||||
ids.append(doc_id)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error indexing {file_path}: {e}")
|
||||
|
||||
# Batch embed and store
|
||||
if documents:
|
||||
embeddings = self.embedder.encode(documents).tolist()
|
||||
self.collection.add(embeddings=embeddings, documents=documents, metadatas=metadatas, ids=ids)
|
||||
|
||||
def _split_code_intelligently(self, content: str, file_path: Path) -> List[Dict[str, Any]]:
|
||||
"""Split code into meaningful chunks."""
|
||||
# Simple line-based splitting for now
|
||||
# Note: Consider using AST for better splitting in future
|
||||
chunks = []
|
||||
lines = content.split("\n")
|
||||
|
||||
# Group into function-sized chunks
|
||||
current_chunk = []
|
||||
current_line = 1
|
||||
|
||||
for i, line in enumerate(lines):
|
||||
current_chunk.append(line)
|
||||
|
||||
# Split on function/class definitions or every 50 lines
|
||||
if len(current_chunk) >= 50 or any(kw in line for kw in ["def ", "function ", "class ", "interface "]):
|
||||
if current_chunk:
|
||||
chunks.append(
|
||||
{
|
||||
"text": "\n".join(current_chunk),
|
||||
"line": current_line,
|
||||
"type": "code",
|
||||
}
|
||||
)
|
||||
current_chunk = []
|
||||
current_line = i + 2
|
||||
|
||||
# Add remaining
|
||||
if current_chunk:
|
||||
chunks.append({"text": "\n".join(current_chunk), "line": current_line, "type": "code"})
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
# Tool registration
|
||||
def create_search_tool(enable_vector_index: bool = False):
|
||||
"""Factory function to create search tool.
|
||||
|
||||
Args:
|
||||
enable_vector_index: Enable vector/semantic search (requires chromadb + sentence-transformers).
|
||||
Default False for fast startup without database dependencies.
|
||||
"""
|
||||
return SearchTool(enable_vector_index=enable_vector_index)
|
||||
def create_search_tool():
|
||||
"""Factory function to create search tool - fast and lightweight."""
|
||||
return SearchTool()
|
||||
|
||||
@@ -124,7 +124,17 @@ class AutoBackgroundExecutor:
|
||||
if wait_task in done:
|
||||
# Process completed within timeout
|
||||
return_code = await wait_task
|
||||
await read_task # Ensure all output is read
|
||||
|
||||
# Give read_task a short time to finish (0.5s max)
|
||||
# This prevents hanging if stdout has issues
|
||||
try:
|
||||
await asyncio.wait_for(read_task, timeout=0.5)
|
||||
except asyncio.TimeoutError:
|
||||
read_task.cancel()
|
||||
try:
|
||||
await read_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# Mark process as completed
|
||||
self.process_manager.mark_completed(process_id, return_code)
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hanzo-mcp"
|
||||
version = "0.9.10"
|
||||
version = "0.9.13"
|
||||
description = "The Zen of Hanzo MCP: One server to rule them all. The ultimate MCP that orchestrates all others."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
|
||||
@@ -196,16 +196,14 @@ async def test_search_pagination():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_vector_opt_in():
|
||||
"""Test that vector search is opt-in and disabled by default."""
|
||||
# Default: vector search disabled
|
||||
async def test_search_no_vector():
|
||||
"""Test that search tool is lightweight without vector/ML dependencies."""
|
||||
# Vector search has been removed - tool should be fast and lightweight
|
||||
search_tool = create_search_tool()
|
||||
assert search_tool._enable_vector_index is False
|
||||
assert search_tool._vector_initialized is False
|
||||
|
||||
# Explicitly enable vector search
|
||||
search_tool_with_vector = create_search_tool(enable_vector_index=True)
|
||||
assert search_tool_with_vector._enable_vector_index is True
|
||||
# No vector-related attributes should exist
|
||||
assert not hasattr(search_tool, '_enable_vector_index')
|
||||
assert not hasattr(search_tool, 'embedder')
|
||||
assert not hasattr(search_tool, 'vector_db')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,244 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Tests for MCP startup performance.
|
||||
|
||||
These tests ensure that MCP server startup remains fast by checking
|
||||
import times. Slow imports cause MCP connections to timeout and hang.
|
||||
|
||||
IMPORTANT: If these tests fail, it means imports have become slow again.
|
||||
The fix is to make heavy imports LAZY - only import when actually needed,
|
||||
not at module load time.
|
||||
|
||||
Common culprits for slow imports:
|
||||
- sentence_transformers (2.5+ seconds)
|
||||
- litellm (1+ second)
|
||||
- hanzo_memory embedding services (3+ seconds)
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
# Maximum allowed import times (in seconds)
|
||||
# These are intentionally generous to avoid flaky tests
|
||||
MAX_MODULE_IMPORT_TIME = 0.5 # hanzo_mcp.tools module
|
||||
MAX_TOTAL_IMPORT_TIME = 1.0 # All core imports combined
|
||||
MAX_CLI_STARTUP_TIME = 3.0 # CLI --help should respond quickly
|
||||
|
||||
|
||||
class TestImportPerformance:
|
||||
"""Test that imports are fast enough for MCP to work."""
|
||||
|
||||
def test_tools_module_import_is_fast(self):
|
||||
"""Test that hanzo_mcp.tools imports quickly.
|
||||
|
||||
This is critical because slow imports cause MCP to hang.
|
||||
The tools module must NOT import heavy dependencies at load time.
|
||||
"""
|
||||
code = '''
|
||||
import time
|
||||
start = time.time()
|
||||
import hanzo_mcp.tools
|
||||
elapsed = time.time() - start
|
||||
print(f"ELAPSED:{elapsed:.3f}")
|
||||
'''
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
assert result.returncode == 0, f"Import failed: {result.stderr}"
|
||||
|
||||
# Extract elapsed time
|
||||
for line in result.stdout.split("\n"):
|
||||
if line.startswith("ELAPSED:"):
|
||||
elapsed = float(line.split(":")[1])
|
||||
break
|
||||
else:
|
||||
pytest.fail(f"Could not find elapsed time in output: {result.stdout}")
|
||||
|
||||
assert elapsed < MAX_MODULE_IMPORT_TIME, (
|
||||
f"hanzo_mcp.tools import took {elapsed:.2f}s (max: {MAX_MODULE_IMPORT_TIME}s). "
|
||||
"This is too slow! Check for heavy imports at module load time. "
|
||||
"Common culprits: sentence_transformers, litellm, hanzo_memory"
|
||||
)
|
||||
|
||||
def test_no_heavy_imports_at_module_level(self):
|
||||
"""Test that heavy packages are NOT imported at module load time."""
|
||||
code = '''
|
||||
import sys
|
||||
# Clear any cached modules
|
||||
for mod in list(sys.modules.keys()):
|
||||
if any(x in mod for x in ['hanzo', 'sentence', 'litellm']):
|
||||
del sys.modules[mod]
|
||||
|
||||
# Import the tools module
|
||||
import hanzo_mcp.tools
|
||||
|
||||
# Check what got imported
|
||||
heavy_modules = []
|
||||
for name in ['sentence_transformers', 'litellm', 'hanzo_memory']:
|
||||
if name in sys.modules:
|
||||
heavy_modules.append(name)
|
||||
|
||||
if heavy_modules:
|
||||
print(f"FAIL:Heavy modules imported at load time: {heavy_modules}")
|
||||
sys.exit(1)
|
||||
else:
|
||||
print("PASS:No heavy modules imported")
|
||||
'''
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
if "FAIL:" in result.stdout:
|
||||
pytest.fail(result.stdout.split("FAIL:")[1].strip())
|
||||
|
||||
assert "PASS:" in result.stdout, f"Unexpected output: {result.stdout}"
|
||||
|
||||
def test_total_import_time(self):
|
||||
"""Test that all core imports combined are fast."""
|
||||
code = '''
|
||||
import time
|
||||
start = time.time()
|
||||
import hanzo_mcp
|
||||
import hanzo_mcp.tools
|
||||
from hanzo_mcp.server import create_server
|
||||
elapsed = time.time() - start
|
||||
print(f"ELAPSED:{elapsed:.3f}")
|
||||
'''
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
assert result.returncode == 0, f"Import failed: {result.stderr}"
|
||||
|
||||
for line in result.stdout.split("\n"):
|
||||
if line.startswith("ELAPSED:"):
|
||||
elapsed = float(line.split(":")[1])
|
||||
break
|
||||
else:
|
||||
pytest.fail(f"Could not find elapsed time in output: {result.stdout}")
|
||||
|
||||
assert elapsed < MAX_TOTAL_IMPORT_TIME, (
|
||||
f"Total imports took {elapsed:.2f}s (max: {MAX_TOTAL_IMPORT_TIME}s). "
|
||||
"Imports are too slow! MCP connections will timeout."
|
||||
)
|
||||
|
||||
|
||||
class TestCLIPerformance:
|
||||
"""Test that CLI starts quickly."""
|
||||
|
||||
def test_cli_help_is_fast(self):
|
||||
"""Test that --help responds quickly."""
|
||||
start = time.time()
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-m", "hanzo_mcp", "--help"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
elapsed = time.time() - start
|
||||
|
||||
assert result.returncode == 0, f"CLI help failed: {result.stderr}"
|
||||
assert elapsed < MAX_CLI_STARTUP_TIME, (
|
||||
f"CLI --help took {elapsed:.2f}s (max: {MAX_CLI_STARTUP_TIME}s). "
|
||||
"This is too slow for MCP connections."
|
||||
)
|
||||
|
||||
def test_cli_does_not_hang(self):
|
||||
"""Test that CLI starts without hanging.
|
||||
|
||||
This was the original issue - MCP tools would hang for 4+ hours
|
||||
because imports were too slow.
|
||||
"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-m", "hanzo_mcp", "--help"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10, # Should respond within 10 seconds max
|
||||
)
|
||||
assert "hanzo" in result.stdout.lower() or "mcp" in result.stdout.lower()
|
||||
except subprocess.TimeoutExpired:
|
||||
pytest.fail(
|
||||
"CLI hung for more than 10 seconds! "
|
||||
"This indicates slow imports causing MCP to hang."
|
||||
)
|
||||
|
||||
|
||||
class TestLazyImportPattern:
|
||||
"""Test that lazy import patterns are correctly implemented."""
|
||||
|
||||
def test_memory_tools_use_lazy_import(self):
|
||||
"""Test that memory tools use lazy imports."""
|
||||
code = '''
|
||||
import sys
|
||||
|
||||
# Import the memory tools module
|
||||
from hanzo_mcp.tools.memory import memory_tools
|
||||
|
||||
# Check that hanzo_memory was NOT imported
|
||||
if 'hanzo_memory' in sys.modules:
|
||||
print("FAIL:hanzo_memory imported at module load")
|
||||
sys.exit(1)
|
||||
print("PASS:hanzo_memory not imported at module load")
|
||||
'''
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
if "FAIL:" in result.stdout:
|
||||
pytest.fail(
|
||||
"memory_tools.py imports hanzo_memory at module load time. "
|
||||
"Use lazy imports with TYPE_CHECKING pattern."
|
||||
)
|
||||
|
||||
def test_search_tool_no_heavy_deps(self):
|
||||
"""Test that search tool does not import heavy ML dependencies.
|
||||
|
||||
Vector search with sentence_transformers has been removed to keep MCP lightweight.
|
||||
If semantic search is needed, use external services (hanzo-node, hanzo desktop).
|
||||
"""
|
||||
code = '''
|
||||
import sys
|
||||
|
||||
# Import the search tool module
|
||||
from hanzo_mcp.tools.search import search_tool
|
||||
|
||||
# Check that sentence_transformers was NOT imported (should be removed entirely)
|
||||
if 'sentence_transformers' in sys.modules:
|
||||
print("FAIL:sentence_transformers imported - should be removed from search_tool")
|
||||
sys.exit(1)
|
||||
print("PASS:no heavy ML dependencies in search_tool")
|
||||
'''
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
if "FAIL:" in result.stdout:
|
||||
pytest.fail(
|
||||
"search_tool.py imports sentence_transformers. "
|
||||
"Vector search has been removed - keep MCP lightweight."
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
Generated
+1
-1
@@ -1003,7 +1003,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "hanzo-mcp"
|
||||
version = "0.9.9"
|
||||
version = "0.9.10"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "bashlex" },
|
||||
|
||||
Reference in New Issue
Block a user