"""Main REPL implementation for Hanzo MCP testing.""" import asyncio import json import os import sys from datetime import datetime from pathlib import Path from typing import Any, Dict, Optional from hanzo_mcp.server import HanzoMCPServer from hanzoai.config import ConfigLoader, RuntimeConfig from hanzoai.protocols import ModelPricing, PermissionMode, PermissionPolicy, UsageTracker from hanzoai.session import ( CompactionConfig, Session as HanzoSession, compact_session, estimate_session_tokens, should_compact, ) from prompt_toolkit import PromptSession from prompt_toolkit.auto_suggest import AutoSuggestFromHistory from prompt_toolkit.completion import WordCompleter from prompt_toolkit.history import FileHistory from rich.console import Console from rich.markdown import Markdown from rich.panel import Panel from rich.table import Table from .llm_client import LLMClient from .tool_executor import ToolExecutor __version__ = "0.1.0" # Default pricing for cost estimates (Claude 3.5 Sonnet tier) _DEFAULT_PRICING = ModelPricing( input_price_per_token=3.0e-6, output_price_per_token=15.0e-6, cache_creation_price_per_token=3.75e-6, cache_read_price_per_token=0.3e-6, ) class HanzoREPL: """Interactive REPL for testing Hanzo MCP tools.""" def __init__(self, config: Optional[Dict[str, Any]] = None): self.console = Console() self.config = config or {} self.mcp_server = None self.llm_client = None self.tool_executor = None self.session = None self.history_file = os.path.expanduser("~/.hanzo_dev_history") # Hierarchical config self.runtime_config: RuntimeConfig = self._load_runtime_config() # Session persistence self._session_path = Path.home() / ".hanzo" / "session.json" self._sessions_dir = Path.home() / ".hanzo" / "sessions" self.hanzo_session = HanzoSession() self.compaction_config = CompactionConfig() # Mode flags self.permission_mode = PermissionMode.Allow self.fast_mode = False # REPL commands self.commands = { "/help": self.show_help, "/tools": self.list_tools, "/exit": self.exit_repl, "/quit": self.exit_repl, "/clear": self.cmd_clear, "/providers": self.list_providers, "/model": self.set_model, "/permissions": self.cmd_permissions, "/context": self.show_context, "/reset": self.reset_context, "/test": self.run_tests, "/status": self.show_status, "/cost": self.show_cost, "/compact": self.run_compact, "/config": self.show_config, "/version": self.show_version, "/login": self.do_login, "/loop": self.do_loop, "/resume": self.cmd_resume, "/memory": self.cmd_memory, "/init": self.cmd_init, "/export": self.cmd_export, "/session": self.cmd_session, "/plan": self.cmd_plan, "/solve": self.cmd_solve, "/code": self.cmd_code, "/auto": self.cmd_auto, "/fast": self.cmd_fast, "/remote-control": self.cmd_remote_control, } def _load_runtime_config(self) -> RuntimeConfig: """Load hierarchical config via ConfigLoader.""" loader = ConfigLoader.default_for(Path.cwd()) return loader.load() async def initialize(self): """Initialize MCP server and LLM client.""" self.console.print("[bold green]Initializing Hanzo REPL...[/bold green]") # Initialize MCP server self.console.print("Loading MCP server...") self.mcp_server = HanzoMCPServer() await self.mcp_server.initialize() # Initialize LLM client self.console.print("Detecting available LLM providers...") self.llm_client = LLMClient() # Apply model from config if present cfg_model = self.runtime_config.get("model") if cfg_model and isinstance(cfg_model, str): try: self.llm_client.set_model(cfg_model) self.console.print(f"[green]Model from config: {cfg_model}[/green]") except ValueError: pass available_providers = self.llm_client.get_available_providers() if not available_providers: self.console.print("[bold red]No LLM API keys detected![/bold red]") self.console.print("Please set one of the following environment variables:") self.console.print("- OPENAI_API_KEY") self.console.print("- ANTHROPIC_API_KEY") self.console.print("- GROQ_API_KEY") self.console.print("- etc.") sys.exit(1) self.console.print( f"[green]Available providers: {', '.join(available_providers)}[/green]" ) # Initialize tool executor with permission policy self.tool_executor = ToolExecutor( self.mcp_server, self.llm_client, permission_policy=PermissionPolicy(default_mode=PermissionMode.Allow), ) # Connect to MCP servers defined in config mcp_servers = self.runtime_config.mcp_servers() for name, cfg in mcp_servers.items(): if hasattr(cfg, "command"): cmd = [cfg.command] + list(cfg.args) count = await self.tool_executor.register_mcp_server( name, cmd, env=cfg.env or None, ) if count: self.console.print(f"[green]MCP {name}: {count} tools[/green]") # Load persisted session if exists if self._session_path.is_file(): try: self.hanzo_session = HanzoSession.load(self._session_path) est = estimate_session_tokens(self.hanzo_session) self.console.print( f"[dim]Restored session: {len(self.hanzo_session.messages)} messages, ~{est:,} tokens[/dim]" ) except Exception: self.hanzo_session = HanzoSession() # Initialize prompt session tool_names = [tool.name for tool in self.mcp_server.tools.values()] completer = WordCompleter( list(self.commands.keys()) + tool_names, ignore_case=True ) self.session = PromptSession( history=FileHistory(self.history_file), auto_suggest=AutoSuggestFromHistory(), completer=completer, ) self.console.print("[bold green]REPL initialized successfully![/bold green]") self.console.print(f"Using model: [cyan]{self.llm_client.current_model}[/cyan]") self.console.print("Type [bold]/help[/bold] for available commands.") async def run(self): """Run the main REPL loop.""" await self.initialize() while True: try: # Get user input user_input = await self.session.prompt_async("hanzo> ", multiline=False) if not user_input.strip(): continue # Check for commands if user_input.startswith("/"): command = user_input.split()[0] if command in self.commands: await self.commands[command](user_input) else: self.console.print(f"[red]Unknown command: {command}[/red]") continue # Process as chat with MCP tools await self.process_chat(user_input) except KeyboardInterrupt: self.console.print("\n[yellow]Use /exit to quit[/yellow]") except EOFError: await self.exit_repl("") except Exception as e: self.console.print(f"[red]Error: {e}[/red]") if self.config.get("debug"): self.console.print_exception() async def process_chat(self, message: str): """Process a chat message with MCP tool support.""" self.console.print() # Send to LLM with available tools try: response = await self.tool_executor.execute_with_tools(message) # Display response if isinstance(response, str): self.console.print(Markdown(response)) else: self.console.print(response) except Exception as e: self.console.print(f"[red]Error processing chat: {e}[/red]") if self.config.get("debug"): self.console.print_exception() async def show_help(self, _): """Show help information.""" help_text = """ # Hanzo REPL Commands ## Session - **/clear** - Clear session history (new session) - **/compact** - Compact session keeping summary - **/resume [path]** - Load saved session from JSON - **/session [list|switch id]** - Manage saved sessions - **/export [file]** - Export conversation to file - **/memory** - Show loaded instruction files - **/init** - Create starter CLAUDE.md in cwd ## Model & Config - **/model [name]** - Show or switch active model - **/permissions [mode]** - Show/switch permission mode - **/config** - Show loaded configuration files - **/fast** - Toggle fast/standard mode - **/providers** - List available LLM providers - **/status** - Model, session info, token usage - **/cost** - Accumulated cost estimate - **/version** - Show version info ## Agents - **/plan [prompt]** - Planning agent - **/solve [prompt]** - Multi-approach solver - **/code [prompt]** - Consensus code review - **/auto [prompt]** - Autonomous execution ## Other - **/login** - Authenticate via PKCE flow - **/loop ** - Run prompt on interval - **/tools** - List available MCP tools - **/context** - Show conversation context - **/reset** - Reset conversation context - **/test** - Run MCP tool tests - **/remote-control** - Start WebSocket remote bridge - **/exit** or **/quit** - Exit the REPL """ self.console.print(Markdown(help_text)) async def list_tools(self, _): """List available MCP tools.""" table = Table(title="Available MCP Tools") table.add_column("Tool", style="cyan", no_wrap=True) table.add_column("Description", style="white") table.add_column("Category", style="green") for tool_name, tool in sorted(self.mcp_server.tools.items()): category = tool.__class__.__module__.split(".")[-1] table.add_row(tool_name, tool.description[:60] + "...", category) self.console.print(table) self.console.print(f"\nTotal tools: [bold]{len(self.mcp_server.tools)}[/bold]") async def list_providers(self, _): """List available LLM providers.""" providers = self.llm_client.get_available_providers() models = self.llm_client.get_available_models() table = Table(title="Available LLM Providers") table.add_column("Provider", style="cyan") table.add_column("Models", style="white") table.add_column("Status", style="green") for provider in providers: provider_models = [m for m in models if m.startswith(provider)] status = "Active" if provider == self.llm_client.current_provider else "" table.add_row( provider, ", ".join(provider_models[:3]) + ("..." if len(provider_models) > 3 else ""), status, ) self.console.print(table) self.console.print( f"\nCurrent model: [bold cyan]{self.llm_client.current_model}[/bold cyan]" ) async def set_model(self, command: str): """Set the LLM model to use.""" parts = command.split(maxsplit=1) if len(parts) < 2: self.console.print("[red]Usage: /model [/red]") self.console.print("Available models:") for model in self.llm_client.get_available_models(): self.console.print(f" - {model}") return model = parts[1] try: self.llm_client.set_model(model) self.console.print(f"[green]Model set to: {model}[/green]") except ValueError as e: self.console.print(f"[red]Error: {e}[/red]") async def show_context(self, _): """Show current conversation context.""" context = self.tool_executor.get_context() if not context: self.console.print("[yellow]No conversation context yet[/yellow]") return self.console.print( Panel( json.dumps(context, indent=2), title="Conversation Context", border_style="blue", ) ) async def reset_context(self, _): """Reset conversation context.""" self.tool_executor.reset_context() self.console.print("[green]Conversation context reset[/green]") async def run_tests(self, _): """Run MCP tool tests.""" self.console.print("[bold]Running MCP tool tests...[/bold]") # Import and run tests from .tests import run_tool_tests await run_tool_tests(self.console, self.mcp_server, self.tool_executor) async def cmd_clear(self, _): """Clear session history and start fresh.""" self.hanzo_session = HanzoSession() self.tool_executor.reset_context() os.system("cls" if os.name == "nt" else "clear") # noqa: S605 self.console.print("[green]Session cleared.[/green]") async def cmd_permissions(self, command: str): """Show or switch permission mode.""" parts = command.split(maxsplit=1) modes = {"read-only": PermissionMode.Deny, "workspace-write": PermissionMode.Ask, "full-access": PermissionMode.Allow} if len(parts) < 2: current = {v: k for k, v in modes.items()}.get(self.permission_mode, "full-access") self.console.print(f"[cyan]Permission mode:[/cyan] {current}") self.console.print(f" Options: {', '.join(modes.keys())}") return name = parts[1].strip() if name not in modes: self.console.print(f"[red]Unknown mode. Use: {', '.join(modes.keys())}[/red]") return self.permission_mode = modes[name] self.tool_executor.permission_policy = PermissionPolicy(default_mode=self.permission_mode) self.console.print(f"[green]Permission mode: {name}[/green]") async def cmd_resume(self, command: str): """Load saved session from JSON file.""" parts = command.split(maxsplit=1) path = Path(parts[1].strip()) if len(parts) > 1 else self._session_path if not path.is_file(): self.console.print(f"[red]No session file at {path}[/red]") return self.hanzo_session = HanzoSession.load(path) est = estimate_session_tokens(self.hanzo_session) self.console.print(f"[green]Loaded: {len(self.hanzo_session.messages)} messages, ~{est:,} tokens[/green]") async def cmd_memory(self, _): """Show loaded instruction files.""" self.console.print("[cyan]Loaded instruction files:[/cyan]") for entry in self.runtime_config.loaded_entries: self.console.print(f" [{entry.source.name}] {entry.path}") if not self.runtime_config.loaded_entries: self.console.print(" [dim]None[/dim]") async def cmd_init(self, _): """Create starter CLAUDE.md in cwd.""" target = Path.cwd() / "CLAUDE.md" if target.exists(): self.console.print(f"[yellow]Already exists: {target}[/yellow]") return target.write_text("# Project Instructions\n\nAdd project-specific instructions here.\n") self.console.print(f"[green]Created {target}[/green]") async def cmd_export(self, command: str): """Export conversation to file.""" parts = command.split(maxsplit=1) outpath = Path(parts[1].strip()) if len(parts) > 1 else Path("hanzo-export.md") lines = [f"# Hanzo Session Export ({datetime.now().isoformat()})\n"] for msg in self.hanzo_session.messages: role = msg.role.value if hasattr(msg.role, "value") else str(msg.role) text = "".join(b.text for b in msg.content if hasattr(b, "text")) lines.append(f"## {role}\n\n{text}\n") outpath.write_text("\n".join(lines)) self.console.print(f"[green]Exported {len(self.hanzo_session.messages)} messages to {outpath}[/green]") async def cmd_session(self, command: str): """List or switch saved sessions.""" parts = command.split(maxsplit=2) self._sessions_dir.mkdir(parents=True, exist_ok=True) if len(parts) < 2 or parts[1] == "list": sessions = sorted(self._sessions_dir.glob("*.json")) if not sessions: self.console.print("[yellow]No saved sessions.[/yellow]") return for s in sessions: self.console.print(f" {s.stem}") return if parts[1] == "switch" and len(parts) > 2: target = self._sessions_dir / f"{parts[2]}.json" if not target.is_file(): self.console.print(f"[red]Session not found: {parts[2]}[/red]") return self.hanzo_session = HanzoSession.load(target) self.console.print(f"[green]Switched to session: {parts[2]}[/green]") return self.console.print("[red]Usage: /session [list|switch ][/red]") async def _agent_chat(self, system_prefix: str, command: str, label: str): """Send a prompt with a system prefix for agent commands.""" parts = command.split(maxsplit=1) if len(parts) < 2: self.console.print(f"[red]Usage: /{label} [/red]") return prompt = f"{system_prefix}\n\nUser request: {parts[1]}" await self.process_chat(prompt) async def cmd_plan(self, command: str): """Planning agent.""" await self._agent_chat("You are a planning agent. Create a detailed plan.", command, "plan") async def cmd_solve(self, command: str): """Multi-approach solver.""" await self._agent_chat("Race multiple approaches. Present the best solution.", command, "solve") async def cmd_code(self, command: str): """Consensus code review.""" await self._agent_chat("Implement this with consensus from multiple review passes.", command, "code") async def cmd_auto(self, command: str): """Autonomous execution.""" await self._agent_chat("Execute this task autonomously. Use tools as needed.", command, "auto") async def cmd_fast(self, _): """Toggle fast/standard mode.""" self.fast_mode = not self.fast_mode state = "fast" if self.fast_mode else "standard" self.console.print(f"[cyan]Mode: {state}[/cyan]") async def cmd_remote_control(self, _): """Start WebSocket server for remote control.""" try: import websockets # noqa: F401 except ImportError: self.console.print("[red]Install websockets: pip install websockets[/red]") return port = 9229 self.console.print(f"[cyan]Remote control bridge: ws://localhost:{port}[/cyan]") self.console.print("[yellow]Send JSON messages: {{\"type\": \"prompt\", \"text\": \"...\"}}[/yellow]") self.console.print("[dim]Press Ctrl+C to stop.[/dim]") async def handler(ws): async for raw in ws: try: msg = json.loads(raw) if msg.get("type") == "prompt": await self.process_chat(msg["text"]) await ws.send(json.dumps({"type": "done"})) except Exception as e: await ws.send(json.dumps({"type": "error", "message": str(e)})) import websockets try: async with websockets.serve(handler, "localhost", port): await asyncio.Future() # run forever except KeyboardInterrupt: self.console.print("[yellow]Remote control stopped.[/yellow]") async def show_status(self, _): """Show model, session info, and token usage.""" tracker = self.tool_executor.usage_tracker if self.tool_executor else UsageTracker() cum = tracker.cumulative_usage() est_tokens = estimate_session_tokens(self.hanzo_session) table = Table(title="Status") table.add_column("Key", style="cyan") table.add_column("Value", style="white") table.add_row("Model", self.llm_client.current_model if self.llm_client else "n/a") table.add_row("Provider", self.llm_client.current_provider if self.llm_client else "n/a") table.add_row("Turns", str(tracker.turns)) table.add_row("Input tokens", f"{cum.input_tokens:,}") table.add_row("Output tokens", f"{cum.output_tokens:,}") table.add_row("Session tokens (est)", f"{est_tokens:,}") table.add_row("Session messages", str(len(self.hanzo_session.messages))) table.add_row("Config files loaded", str(len(self.runtime_config.loaded_entries))) self.console.print(table) async def show_cost(self, _): """Show accumulated cost estimate.""" tracker = self.tool_executor.usage_tracker if self.tool_executor else UsageTracker() cum = tracker.cumulative_usage() cost = _DEFAULT_PRICING.cost(cum) self.console.print( f"[cyan]Cost estimate:[/cyan] ${cost:.6f} " f"({cum.input_tokens:,} in / {cum.output_tokens:,} out, {tracker.turns} turns)" ) async def run_compact(self, _): """Compact the session to reclaim context.""" est = estimate_session_tokens(self.hanzo_session) result = compact_session(self.hanzo_session, self.compaction_config) if result.removed_message_count == 0: self.console.print("[yellow]Session too small to compact.[/yellow]") return self.hanzo_session = result.compacted_session new_est = estimate_session_tokens(self.hanzo_session) self.console.print( f"[green]Compacted: removed {result.removed_message_count} messages, " f"{est:,} -> {new_est:,} est tokens.[/green]" ) async def show_config(self, _): """Show loaded configuration.""" if not self.runtime_config.loaded_entries: self.console.print("[yellow]No config files loaded.[/yellow]") return for entry in self.runtime_config.loaded_entries: self.console.print(f" [{entry.source.name}] {entry.path}") mcp_servers = self.runtime_config.mcp_servers() if mcp_servers: self.console.print(f"\n MCP servers: {', '.join(mcp_servers.keys())}") async def show_version(self, _): """Show version info.""" from hanzoai._version import __version__ as hanzoai_version self.console.print(f" hanzo-dev {__version__}") self.console.print(f" hanzoai {hanzoai_version}") async def do_login(self, _): """Login via PKCE flow.""" from hanzoai.auth import HanzoAuth auth = HanzoAuth() oauth_cfg = self.runtime_config.oauth() kwargs: Dict[str, Any] = {} if oauth_cfg: kwargs["authorize_url"] = oauth_cfg.authorize_url kwargs["token_url"] = oauth_cfg.token_url kwargs["client_id"] = oauth_cfg.client_id if oauth_cfg.scopes: kwargs["scopes"] = oauth_cfg.scopes if oauth_cfg.callback_port: kwargs["redirect_port"] = oauth_cfg.callback_port self.console.print("[yellow]Starting PKCE login flow...[/yellow]") try: token_set = await auth.login_with_pkce(**kwargs) self.console.print(f"[green]Logged in. Scopes: {', '.join(token_set.scopes)}[/green]") except Exception as e: self.console.print(f"[red]Login failed: {e}[/red]") async def do_loop(self, command: str): """Run a prompt on a recurring interval. Usage: /loop """ parts = command.split(maxsplit=2) if len(parts) < 3: self.console.print("[red]Usage: /loop [/red]") return try: interval = int(parts[1]) except ValueError: self.console.print("[red]Interval must be an integer (seconds).[/red]") return prompt = parts[2] self.console.print(f"[yellow]Looping every {interval}s: {prompt}[/yellow]") try: while True: await self.process_chat(prompt) await asyncio.sleep(interval) except KeyboardInterrupt: self.console.print("[yellow]Loop stopped.[/yellow]") async def exit_repl(self, _): """Exit the REPL, saving session.""" # Persist session self._session_path.parent.mkdir(parents=True, exist_ok=True) self.hanzo_session.save(self._session_path) self.console.print("[yellow]Goodbye![/yellow]") sys.exit(0) async def main(): """Main entry point.""" repl = HanzoREPL() await repl.run() if __name__ == "__main__": asyncio.run(main())