jarvis-memory/skills/task-queue/scripts/heartbeat_worker.py

#!/usr/bin/env python3
"""
Heartbeat worker - GPT-powered task execution.
Sends tasks to Ollama for command generation, executes via SSH.
"""

import redis
import json
import time
import os
import sys
import subprocess
import requests
from datetime import datetime

REDIS_HOST = os.environ.get("REDIS_HOST", "127.0.0.1")
REDIS_PORT = int(os.environ.get("REDIS_PORT", 6379))
REDIS_PASSWORD = os.environ.get("REDIS_PASSWORD", None)
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
TASK_LLM_MODEL = os.environ.get("TASK_LLM_MODEL", "kimi-k2.5:cloud")
DEFAULT_TARGET_HOST = os.environ.get("TASK_SSH_HOST", "")
DEFAULT_SSH_USER = os.environ.get("TASK_SSH_USER", "")
DEFAULT_SUDO_PASS = os.environ.get("TASK_SUDO_PASS", "")

def get_redis():
    return redis.Redis(
        host=REDIS_HOST,
        port=REDIS_PORT,
        password=REDIS_PASSWORD,
        decode_responses=True
    )

def generate_task_id():
    return f"task_{int(time.time())}_{os.urandom(4).hex()}"

def check_active_task(r):
    """Check if there's already an active task."""
    active = r.lrange("tasks:active", 0, -1)
    if active:
        task_id = active[0]
        task = r.hgetall(f"task:{task_id}")
        started_at = int(task.get("started_at", 0))
        elapsed = time.time() - started_at
        print(f"[BUSY] Task {task_id} active for {elapsed:.0f}s")
        return True
    return False

def get_pending_task(r):
    """Pop a task from pending queue."""
    task_id = r.rpop("tasks:pending")
    if task_id:
        return task_id
    return None

def clean_json_content(content):
    """Strip markdown code blocks if present."""
    cleaned = content.strip()
    if cleaned.startswith("```json"):
        cleaned = cleaned[7:]
    elif cleaned.startswith("```"):
        cleaned = cleaned[3:]
    if cleaned.endswith("```"):
        cleaned = cleaned[:-3]
    return cleaned.strip()

def ask_gpt_for_commands(task_description, target_host=None, ssh_user=None, sudo_pass=None):
    """
    Send task to Ollama/GPT to generate SSH commands.
    Returns dict with commands, expected results, and explanation.
    """
    target_host = target_host or DEFAULT_TARGET_HOST
    ssh_user = ssh_user or DEFAULT_SSH_USER
    sudo_pass = sudo_pass if sudo_pass is not None else DEFAULT_SUDO_PASS

    if not target_host or not ssh_user:
        raise ValueError("TASK_SSH_HOST and TASK_SSH_USER must be set (or passed explicitly)")

    sudo_line = (
        f"Sudo password: {sudo_pass}"
        if sudo_pass
        else "Sudo password: (not provided; avoid sudo unless absolutely necessary)"
    )

    system_prompt = f"""You have SSH access to {ssh_user}@{target_host}
{sudo_line}

Your job is to generate shell commands to complete the given task.
Respond ONLY with valid JSON in this format:
{{
  "commands": [
    "ssh -t {ssh_user}@{target_host} 'sudo apt update'",
    "ssh -t {ssh_user}@{target_host} 'sudo apt install -y docker.io'"
  ],
  "expected_results": [
    "apt updated successfully",
    "docker installed and running"
  ],
  "explanation": "Updating packages and installing Docker"
}}

Rules:
- Commands should use ssh -t (allocates TTY for sudo) to execute on the remote host
- Use sudo only when needed
- Keep commands safe and idempotent where possible
- If task is unclear, ask for clarification in explanation

For Docker-related tasks:
- Search Docker Hub for official images (docker.io/library/ or verified publishers)
- Prefer latest stable versions
- Use official images over community when available
- Verify image exists before trying to pull
- Map volumes as specified in the task (e.g., -v /root/html:/usr/share/nginx/html)
"""

    user_prompt = f"Task: {task_description}\n\nGenerate the commands to complete this task."

    try:
        response = requests.post(
            f"{OLLAMA_URL}/api/chat",
            json={
                "model": TASK_LLM_MODEL,
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                "stream": False,
                "format": "json"
            },
            timeout=120
        )
        response.raise_for_status()
        
        result = response.json()
        content = result.get("message", {}).get("content", "{}")
        
        # Parse the JSON response
        try:
            cleaned = clean_json_content(content)
            gpt_plan = json.loads(cleaned)
            return gpt_plan
        except json.JSONDecodeError:
            # If GPT didn't return valid JSON, wrap the raw response
            return {
                "commands": [],
                "expected_results": [],
                "explanation": f"GPT response: {content[:200]}",
                "parse_error": "GPT did not return valid JSON"
            }
            
    except Exception as e:
        return {
            "commands": [],
            "expected_results": [],
            "explanation": f"Failed to get commands from GPT: {e}",
            "error": str(e)
        }

def execute_ssh_command_with_sudo(command, sudo_pass, timeout=300):
    """
    Execute an SSH command with sudo password handling.
    Uses -t flag for TTY allocation and handles sudo password prompt.
    """
    try:
        # Ensure command has -t flag for TTY
        if not "-t" in command and command.startswith("ssh "):
            command = command.replace("ssh ", "ssh -t ", 1)
        
        # Use expect-like approach with subprocess
        # Send password when prompted
        import pty
        import select
        import termios
        import tty
        
        master_fd, slave_fd = pty.openpty()
        
        process = subprocess.Popen(
            command,
            shell=True,
            stdin=slave_fd,
            stdout=slave_fd,
            stderr=slave_fd,
            preexec_fn=os.setsid
        )
        
        os.close(slave_fd)
        
        output = []
        password_sent = False
        start_time = time.time()
        
        while process.poll() is None:
            if time.time() - start_time > timeout:
                process.kill()
                return {
                    "success": False,
                    "stdout": "".join(output),
                    "stderr": "Command timed out",
                    "exit_code": -1
                }
            
            ready, _, _ = select.select([master_fd], [], [], 0.1)
            if ready:
                try:
                    data = os.read(master_fd, 1024).decode()
                    output.append(data)
                    
                    # Check for sudo password prompt
                    if "password:" in data.lower() or "password for" in data.lower():
                        if not password_sent:
                            os.write(master_fd, (sudo_pass + "\n").encode())
                            password_sent = True
                            time.sleep(0.5)
                except OSError:
                    break
        
        os.close(master_fd)
        
        stdout = "".join(output)
        return {
            "success": process.returncode == 0,
            "stdout": stdout,
            "stderr": "" if process.returncode == 0 else stdout,
            "exit_code": process.returncode
        }
        
    except Exception as e:
        return {
            "success": False,
            "stdout": "",
            "stderr": str(e),
            "exit_code": -1
        }

def execute_ssh_command_simple(command, timeout=300):
    """
    Execute an SSH command without sudo (simple version).
    """
    try:
        result = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=timeout
        )
        return {
            "success": result.returncode == 0,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "exit_code": result.returncode
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "stdout": "",
            "stderr": "Command timed out",
            "exit_code": -1
        }
    except Exception as e:
        return {
            "success": False,
            "stdout": "",
            "stderr": str(e),
            "exit_code": -1
        }

def execute_task_with_gpt(task):
    """
    Execute task using GPT to generate commands, then run via SSH.
    """
    task_description = task.get("description", "No description")
    target_host = task.get("target_host", "10.0.0.38")
    ssh_user = task.get("ssh_user", "n8n")
    sudo_pass = task.get("sudo_pass", "passw0rd")
    
    print(f"[GPT] Generating commands for: {task_description}")
    
    # Get commands from GPT
    gpt_plan = ask_gpt_for_commands(task_description, target_host, ssh_user, sudo_pass)
    
    if not gpt_plan.get("commands"):
        comments = f"GPT failed to generate commands: {gpt_plan.get('explanation', 'Unknown error')}"
        return {
            "success": False,
            "gpt_plan": gpt_plan,
            "execution_results": [],
            "comments": comments
        }
    
    print(f"[GPT] Plan: {gpt_plan.get('explanation', 'No explanation')}")
    print(f"[EXEC] Running {len(gpt_plan['commands'])} commands...")
    
    # Execute each command
    execution_results = []
    any_failed = False
    
    for i, cmd in enumerate(gpt_plan["commands"]):
        print(f"[CMD {i+1}] {cmd[:80]}...")
        
        # Check if command uses sudo
        if "sudo" in cmd.lower():
            result = execute_ssh_command_with_sudo(cmd, sudo_pass)
        else:
            result = execute_ssh_command_simple(cmd)
            
        execution_results.append({
            "command": cmd,
            "result": result
        })
        
        if not result["success"]:
            any_failed = True
            print(f"[FAIL] Exit code {result['exit_code']}: {result['stderr'][:100]}")
        else:
            print(f"[OK] Success")
    
    # Build comments field
    if any_failed:
        failed_cmds = [r for r in execution_results if not r["result"]["success"]]
        comments = f"ERRORS ({len(failed_cmds)} failed):\n"
        for r in failed_cmds:
            comments += f"- Command: {r['command'][:60]}...\n"
            comments += f"  Error: {r['result']['stderr'][:200]}\n"
    else:
        comments = "OK"
    
    return {
        "success": not any_failed,
        "gpt_plan": gpt_plan,
        "execution_results": execution_results,
        "comments": comments
    }

def execute_simple_task(task):
    """
    Execute simple tasks (notify, command) without GPT.
    """
    task_type = task.get("type", "default")
    description = task.get("description", "No description")
    sudo_pass = task.get("sudo_pass", "passw0rd")
    
    if task_type == "notify":
        # For now, just log it (messaging handled elsewhere)
        return {
            "success": True,
            "result": f"Notification: {task.get('message', description)}",
            "comments": "OK"
        }
    
    elif task_type == "command":
        # Execute shell command directly
        command = task.get("command", "")
        if command:
            if "sudo" in command.lower():
                result = execute_ssh_command_with_sudo(command, sudo_pass)
            else:
                result = execute_ssh_command_simple(command)
            comments = "OK" if result["success"] else f"Error: {result['stderr'][:500]}"
            return {
                "success": result["success"],
                "result": result["stdout"][:500],
                "comments": comments
            }
        else:
            return {
                "success": False,
                "result": "No command specified",
                "comments": "ERROR: No command provided"
            }
    
    else:
        # Default: use GPT
        return execute_task_with_gpt(task)

def mark_completed(r, task_id, result_data):
    """Mark task as completed with full result data."""
    r.hset(f"task:{task_id}", mapping={
        "status": "completed" if result_data["success"] else "failed",
        "completed_at": str(int(time.time())),
        "result": json.dumps(result_data.get("result", "")),
        "comments": result_data.get("comments", "")
    })
    r.lrem("tasks:active", 0, task_id)
    r.lpush("tasks:completed", task_id)
    
    status = "DONE" if result_data["success"] else "FAILED"
    print(f"[{status}] {task_id}")
    if result_data.get("comments") and result_data["comments"] != "OK":
        print(f"[COMMENTS] {result_data['comments'][:200]}")

def mark_failed(r, task_id, error):
    """Mark task as failed."""
    r.hset(f"task:{task_id}", mapping={
        "status": "failed",
        "completed_at": str(int(time.time())),
        "result": f"Error: {error}",
        "comments": f"Worker error: {error}"
    })
    r.lrem("tasks:active", 0, task_id)
    r.lpush("tasks:completed", task_id)
    print(f"[FAILED] {task_id}: {error}")

def main():
    r = get_redis()
    
    # Check if already busy
    if check_active_task(r):
        sys.exit(0)
    
    # Get next pending task
    task_id = get_pending_task(r)
    if not task_id:
        print("[IDLE] No pending tasks")
        sys.exit(0)
    
    # Load task details
    task = r.hgetall(f"task:{task_id}")
    if not task:
        print(f"[ERROR] Task {task_id} not found")
        sys.exit(1)
    
    # Move to active
    r.hset(f"task:{task_id}", mapping={
        "status": "active",
        "started_at": str(int(time.time()))
    })
    r.lpush("tasks:active", task_id)
    
    print(f"[START] {task_id}: {task.get('description', 'No description')}")
    
    try:
        # Execute the task
        result_data = execute_simple_task(task)
        mark_completed(r, task_id, result_data)
        print(f"[WAKE] Task complete - check comments field for status")
        
    except Exception as e:
        mark_failed(r, task_id, str(e))
        sys.exit(1)

if __name__ == "__main__":
    main()