mnemo_cards/tools/agent/cursor_cli_wrapper.py

"""
Cursor CLI wrapper with stream-json parsing and progress tracking.
"""
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Optional, Dict, Any, List, Callable
from dataclasses import dataclass
from enum import Enum


class CursorResultStatus(Enum):
    """Status of cursor CLI execution."""
    SUCCESS = "success"
    FAILURE = "failure"
    TIMEOUT = "timeout"
    ERROR = "error"


@dataclass
class CursorResult:
    """Result of cursor CLI execution."""
    status: CursorResultStatus
    output: str
    error: Optional[str]
    duration_ms: Optional[int]
    files_modified: List[str]
    files_created: List[str]
    files_read: List[str]
    tool_calls: int
    

class CursorCLI:
    """Wrapper for Cursor CLI with stream-json parsing."""
    
    def __init__(self, 
                 project_root: Path, 
                 api_key: str,
                 model: str = "auto",
                 verbose: bool = True):
        self.project_root = project_root
        self.api_key = api_key
        self.model = model
        self.verbose = verbose
        self.cursor_agent_path = self._find_cursor_agent()
    
    def _find_cursor_agent(self) -> str:
        """Find cursor-agent executable."""
        # Try to find in PATH first
        cursor_agent = shutil.which("cursor-agent")
        if cursor_agent:
            return cursor_agent
        
        # Try common installation locations (new default location first)
        home = os.path.expanduser("~")
        possible_paths = [
            os.path.join(home, ".local", "bin", "cursor-agent"),  # New default location
            os.path.join(home, ".cursor", "bin", "cursor-agent"),  # Old location
            "/usr/local/bin/cursor-agent",
            "/usr/bin/cursor-agent",
        ]
        
        for path in possible_paths:
            if os.path.isfile(path) and os.access(path, os.X_OK):
                if self.verbose:
                    print(f"✅ Found cursor-agent at: {path}")
                return path
        
        # If not found, return "cursor-agent" and let it fail with a clear error
        if self.verbose:
            print("⚠️  cursor-agent not found in PATH or common locations")
            print(f"   PATH: {os.environ.get('PATH', 'not set')}")
            print(f"   Home: {home}")
            print(f"   Checked locations: {possible_paths}")
            print("   Trying 'cursor-agent' anyway...")
        
        return "cursor-agent"
    
    def run_agent(self,
                  task_description: str,
                  force: bool = True,
                  max_iterations: int = 5,
                  progress_callback: Optional[Callable[[str, Any], None]] = None) -> CursorResult:
        """
        Run cursor agent on a task.
        
        Args:
            task_description: Description of the task for the agent
            force: Allow file modifications
            max_iterations: Maximum number of agent iterations
            progress_callback: Callback function(event_type, data) for progress updates
        
        Returns:
            CursorResult with execution details
        """
        # Build command using found path
        cmd = [
            self.cursor_agent_path,
            "-p",  # Print mode (non-interactive)
            "--output-format", "stream-json",
            "--stream-partial-output",
        ]
        
        if force:
            cmd.append("--force")
        
        # Add task description
        cmd.append(task_description)
        
        # Set environment with PATH including bash location
        current_path = os.environ.get("PATH", "")
        # Ensure common bash locations are in PATH
        bash_paths = [
            "/usr/bin",
            "/bin",
            os.path.expanduser("~/.local/bin"),
            os.path.expanduser("~/.cursor/bin"),
        ]
        
        # Add bash paths if they exist
        additional_paths = ":".join([p for p in bash_paths if os.path.exists(p)])
        new_path = ":".join([additional_paths, current_path]) if additional_paths else current_path
        # Prepare environment variables
        env = {
            "CURSOR_API_KEY": self.api_key,
            "PATH": new_path,
        }
        
        # Prepare command
        cmd = [self.cursor_agent_path]
        
        # Add model flag if provided
        if self.model:
             cmd.extend(["--model", self.model])
             
        # Add prompt
        cmd.append(task_description)
        
        # Add options
        cmd.extend(["--output-format", "stream-json"])
        cmd.append("--stream-partial-output")
        
        if force:
            cmd.append("--force")
        
        # Add print mode (non-interactive)
        cmd.append("-p")
        
        if self.verbose:
            print(f"🚀 Running cursor-agent in {self.project_root}")
            print(f"📝 Task: {task_description[:100]}...")
            print(f"🔧 Using cursor-agent at: {self.cursor_agent_path}")
        
        # Verify cursor-agent exists
        if self.cursor_agent_path == "cursor-agent":
            # Try one more time to find it
            cursor_agent = shutil.which("cursor-agent")
            if not cursor_agent:
                error_msg = (
                    f"❌ cursor-agent not found!\n"
                    f"   Please ensure Cursor CLI is installed.\n"
                    f"   Installation: curl https://cursor.com/install -fsS | bash\n"
                    f"   Expected locations:\n"
                    f"     - $HOME/.local/bin/cursor-agent (new default)\n"
                    f"     - $HOME/.cursor/bin/cursor-agent (old location)\n"
                    f"   Current PATH: {os.environ.get('PATH', 'not set')}\n"
                    f"   Current HOME: {os.environ.get('HOME', 'not set')}"
                )
                return CursorResult(
                    status=CursorResultStatus.ERROR,
                    output="",
                    error=error_msg,
                    duration_ms=None,
                    files_modified=[],
                    files_created=[],
                    files_read=[],
                    tool_calls=0
                )
            self.cursor_agent_path = cursor_agent
            if self.verbose:
                print(f"✅ Found cursor-agent at: {self.cursor_agent_path}")
        
        # Track execution metrics
        accumulated_text = ""
        tool_count = 0
        files_modified = []
        files_created = []
        files_read = []
        duration_ms = None
        error_message = None
        
        try:
            # Run cursor-agent with streaming output
            process = subprocess.Popen(
                cmd,
                cwd=self.project_root,
                env={**subprocess.os.environ, **env},
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1
            )
            
            # Process streaming output line by line
            for line in process.stdout:
                line = line.strip()
                if not line:
                    continue
                
                try:
                    event = json.loads(line)
                    event_type = event.get("type")
                    subtype = event.get("subtype")
                    
                    # Call progress callback
                    if progress_callback:
                        progress_callback(event_type, event)
                    
                    # Process different event types
                    if event_type == "system":
                        self._handle_system_event(event, subtype)
                    
                    elif event_type == "assistant":
                        accumulated_text += self._handle_assistant_event(event)
                    
                    elif event_type == "tool_call":
                        tool_result = self._handle_tool_call_event(
                            event, subtype, files_modified, files_created, files_read
                        )
                        if tool_result:
                            tool_count += 1
                    
                    elif event_type == "result":
                        duration_ms = event.get("duration_ms")
                        if self.verbose:
                            self._print_summary(
                                duration_ms, tool_count, 
                                len(accumulated_text),
                                files_created, files_modified
                            )
                
                except json.JSONDecodeError:
                    # Not JSON, might be error message
                    if self.verbose:
                        print(f"⚠️  {line}")
                    error_message = line
            
            # Wait for process to complete
            return_code = process.wait()
            
            # Check stderr
            stderr = process.stderr.read()
            if stderr:
                error_message = stderr
                if self.verbose:
                    print(f"❌ Error: {stderr}")
            
            # Determine status
            if return_code == 0:
                status = CursorResultStatus.SUCCESS
            else:
                status = CursorResultStatus.FAILURE
            
            return CursorResult(
                status=status,
                output=accumulated_text,
                error=error_message,
                duration_ms=duration_ms,
                files_modified=files_modified,
                files_created=files_created,
                files_read=files_read,
                tool_calls=tool_count
            )
        
        except subprocess.TimeoutExpired:
            return CursorResult(
                status=CursorResultStatus.TIMEOUT,
                output=accumulated_text,
                error="Cursor CLI execution timeout",
                duration_ms=None,
                files_modified=files_modified,
                files_created=files_created,
                files_read=files_read,
                tool_calls=tool_count
            )
        
        except Exception as e:
            return CursorResult(
                status=CursorResultStatus.ERROR,
                output=accumulated_text,
                error=str(e),
                duration_ms=None,
                files_modified=files_modified,
                files_created=files_created,
                files_read=files_read,
                tool_calls=tool_count
            )
    
    def _handle_system_event(self, event: Dict[str, Any], subtype: Optional[str]) -> None:
        """Handle system event."""
        if subtype == "init" and self.verbose:
            model = event.get("model", "unknown")
            print(f"🤖 Using model: {model}")
    
    def _handle_assistant_event(self, event: Dict[str, Any]) -> str:
        """Handle assistant event and return text delta."""
        content = event.get("message", {}).get("content", [])
        if content:
            text = content[0].get("text", "")
            if text and self.verbose:
                # Show progress indicator
                sys.stdout.write(".")
                sys.stdout.flush()
            return text
        return ""
    
    def _handle_tool_call_event(self,
                                event: Dict[str, Any],
                                subtype: Optional[str],
                                files_modified: List[str],
                                files_created: List[str],
                                files_read: List[str]) -> bool:
        """Handle tool call event. Returns True if tool completed."""
        tool_call = event.get("tool_call", {})
        
        if subtype == "started":
            # Tool started
            if "writeToolCall" in tool_call:
                path = tool_call["writeToolCall"].get("args", {}).get("path", "unknown")
                if self.verbose:
                    print(f"\n🔧 Writing: {path}")
            
            elif "readToolCall" in tool_call:
                path = tool_call["readToolCall"].get("args", {}).get("path", "unknown")
                files_read.append(path)
                if self.verbose:
                    print(f"\n📖 Reading: {path}")
            
            return False
        
        elif subtype == "completed":
            # Tool completed
            if "writeToolCall" in tool_call:
                result = tool_call["writeToolCall"].get("result", {})
                if "success" in result:
                    success = result["success"]
                    path = tool_call["writeToolCall"].get("args", {}).get("path", "unknown")
                    lines_created = success.get("linesCreated", 0)
                    lines_modified = success.get("linesModified", 0)
                    
                    if lines_created > 0:
                        files_created.append(path)
                        if self.verbose:
                            print(f"   ✅ Created {lines_created} lines")
                    elif lines_modified > 0:
                        files_modified.append(path)
                        if self.verbose:
                            print(f"   ✅ Modified {lines_modified} lines")
                    else:
                        # File was written but not actually changed
                        if self.verbose:
                            print(f"   ⚠️  Written but no changes detected (file unchanged)")
            
            elif "readToolCall" in tool_call:
                result = tool_call["readToolCall"].get("result", {})
                if "success" in result and self.verbose:
                    lines = result["success"].get("totalLines", 0)
                    print(f"   ✅ Read {lines} lines")
            
            return True
        
        return False
    
    def _print_summary(self, 
                      duration_ms: Optional[int],
                      tool_count: int,
                      chars_generated: int,
                      files_created: List[str],
                      files_modified: List[str]) -> None:
        """Print execution summary."""
        print(f"\n\n🎯 Completed in {duration_ms}ms")
        print(f"📊 Stats: {tool_count} tool calls, {chars_generated} chars generated")
        if files_created:
            print(f"📝 Created {len(files_created)} file(s)")
        if files_modified:
            print(f"✏️  Modified {len(files_modified)} file(s)")


def test_cursor_cli():
    """Test cursor CLI wrapper."""
    import os
    from pathlib import Path
    
    api_key = os.getenv("CURSOR_API_KEY")
    if not api_key:
        print("❌ CURSOR_API_KEY not set")
        return
    
    project_root = Path.cwd()
    
    cli = CursorCLI(project_root, api_key)
    
    task = "List all Python files in tools/agent directory"
    
    def progress_callback(event_type: str, event: Dict[str, Any]):
        """Example progress callback."""
        if event_type == "tool_call" and event.get("subtype") == "started":
            print(f"🔄 Tool started...")
    
    result = cli.run_agent(
        task_description=task,
        force=False,
        progress_callback=progress_callback
    )
    
    print(f"\n📋 Result: {result.status}")
    print(f"📄 Output: {result.output[:200]}...")


if __name__ == "__main__":
    test_cursor_cli()