nas-tools/dashboard/backend/services/agent_executor.py

"""
Agent Executor Service - Executes agent tasks and manages their lifecycle
"""

import asyncio
import json
import logging
import os
from datetime import datetime
from typing import Any

import httpx

from db import opc_db
from services import email_service
from services.agents.base_agent import BaseAgent

logger = logging.getLogger(__name__)

# Telegram notification settings
TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
TELEGRAM_CHAT_ID = os.getenv("TELEGRAM_CHAT_ID")


class AgentExecutor:
    """Manages agent execution lifecycle"""

    def __init__(self):
        self.agents: dict[str, BaseAgent] = {}
        self.running = False

    async def initialize(self):
        """Load agents from database"""
        agents_data = await opc_db.get_agents(enabled_only=True)

        for agent_data in agents_data:
            agent = BaseAgent(
                agent_id=agent_data["id"],
                name=agent_data["name"],
                role=agent_data["role"],
                system_prompt=agent_data["system_prompt"],
                config=agent_data["config"],
            )
            self.agents[agent_data["id"]] = agent
            logger.info(f"Loaded agent: {agent.name} ({agent.agent_id})")

    async def start(self):
        """Start the executor service"""
        self.running = True
        print("=== Agent Executor Loop Started ===")
        logger.info("Agent executor service started")

        while self.running:
            try:
                print(f"[Executor] Checking for pending executions... (running={self.running})")
                await self._process_pending_executions()
            except Exception as e:
                print(f"[Executor] Error in executor loop: {e}")
                logger.error(f"Error in executor loop: {e}")
                import traceback

                traceback.print_exc()

            await asyncio.sleep(5)  # Check every 5 seconds

    async def stop(self):
        """Stop the executor service"""
        self.running = False
        logger.info("Agent executor service stopped")

    async def _process_pending_executions(self):
        """Process pending agent executions"""
        # Process both pending and approved executions
        pending_executions = await opc_db.get_agent_executions(status="pending", limit=10)
        approved_executions = await opc_db.get_agent_executions(status="approved", limit=10)

        executions = pending_executions + approved_executions
        print(f"[Executor] Found {len(pending_executions)} pending and {len(approved_executions)} approved executions")

        for execution in executions:
            try:
                print(
                    f"[Executor] Processing execution {execution['id']} for agent {execution['agent_id']} (status: {execution['status']})"
                )

                if execution["status"] == "approved":
                    # Resume approved execution
                    await self._resume_approved_execution(execution)
                else:
                    # Execute new pending execution
                    await self._execute_agent(execution)
            except Exception as e:
                print(f"[Executor] Failed to execute agent {execution['agent_id']}: {e}")
                logger.error(f"Failed to execute agent {execution['agent_id']}: {e}")
                await self._mark_execution_failed(execution["id"], str(e))

    async def _execute_agent(self, execution: dict[str, Any]):
        """Execute a single agent task"""
        agent_id = execution["agent_id"]
        task_id = execution["task_id"]

        # Get agent
        agent = self.agents.get(agent_id)
        if not agent:
            raise ValueError(f"Agent {agent_id} not found")

        # Mark as running
        await self._update_execution_status(execution["id"], "running")

        # Get task and context
        task = await opc_db.get_task_by_id(task_id)
        if not task:
            raise ValueError(f"Task {task_id} not found")

        # Build context
        context = await self._build_context(task)

        # Execute agent
        logger.info(f"Executing agent {agent.name} on task {task['title']}")
        result = await agent.execute(task, context)

        # Check if requires approval
        if execution.get("requires_approval") or result.get("requires_approval"):
            # Mark as pending approval
            await self._mark_pending_approval(execution["id"], result)
            await self._send_approval_notification(agent, task, result)
        else:
            # Execute actions
            await self._execute_actions(task, result["actions"])

            # Mark as completed
            await self._mark_execution_completed(execution["id"], result)

            # Send completion notification
            await self._send_completion_notification(agent, task, result)

    async def _resume_approved_execution(self, execution: dict[str, Any]):
        """Resume an approved execution and execute its actions"""
        agent_id = execution["agent_id"]
        task_id = execution["task_id"]

        # Get agent
        agent = self.agents.get(agent_id)
        if not agent:
            raise ValueError(f"Agent {agent_id} not found")

        # Get task
        task = await opc_db.get_task_by_id(task_id)
        if not task:
            raise ValueError(f"Task {task_id} not found")

        # Mark as running
        await self._update_execution_status(execution["id"], "running")

        # Get proposed actions from execution
        result = execution.get("output_result", {})
        actions = execution.get("actions_proposed", [])

        logger.info(f"Resuming approved execution for agent {agent.name} on task {task['title']}")

        # Execute approved actions
        await self._execute_actions(task, actions)

        # Mark as completed
        await self._mark_execution_completed(execution["id"], result)

        # Send completion notification
        await self._send_completion_notification(agent, task, result)

    async def _build_context(self, task: dict[str, Any]) -> dict[str, Any]:
        """Build context for agent execution"""
        context = {"task": task, "timestamp": datetime.utcnow().isoformat()}

        # Get related tasks from same project
        if task.get("project_id"):
            related_tasks = await opc_db.get_tasks(project_id=task["project_id"], limit=20)
            context["related_tasks"] = related_tasks

        # Get task history
        history = await opc_db.get_task_history(task["id"])
        context["history"] = history[:10]  # Last 10 events

        return context

    async def _execute_actions(self, task: dict[str, Any], actions: list):
        """Execute proposed actions"""
        for action in actions:
            action_type = action.get("type")

            try:
                if action_type == "create_subtask":
                    await self._action_create_subtask(task, action)
                elif action_type == "update_task_status":
                    await self._action_update_status(task, action)
                elif action_type == "send_notification":
                    await self._action_send_notification(action)
                elif action_type == "add_comment":
                    await self._action_add_comment(task, action)
                elif action_type == "error":
                    # Handle error action - log and fail execution
                    error_msg = action.get("message", "Unknown error")
                    logger.error(f"Agent reported error: {error_msg}")
                    raise Exception(error_msg)
                else:
                    logger.warning(f"Unknown action type: {action_type}")
            except Exception as e:
                logger.error(f"Failed to execute action {action_type}: {e}")
                raise

    async def _action_create_subtask(self, parent_task: dict[str, Any], action: dict[str, Any]):
        """Create a subtask"""
        await opc_db.create_task(
            title=action.get("title", "Subtask"),
            description=action.get("description", ""),
            status="parking_lot",
            priority=parent_task.get("priority", "medium"),
            project_id=parent_task.get("project_id"),
            tags=parent_task.get("tags", []),
            actor="agent",
        )
        logger.info(f"Created subtask: {action.get('title')}")

    async def _action_update_status(self, task: dict[str, Any], action: dict[str, Any]):
        """Update task status"""
        new_status = action.get("status")
        if new_status:
            await opc_db.update_task(task_id=task["id"], updates={"status": new_status}, actor="agent")
            logger.info(f"Updated task {task['id']} status to {new_status}")

    async def _action_send_notification(self, action: dict[str, Any]):
        """Send notification"""
        message = action.get("message", "")
        if message and TELEGRAM_BOT_TOKEN and TELEGRAM_CHAT_ID:
            await self._send_telegram(message)

    async def _action_add_comment(self, task: dict[str, Any], action: dict[str, Any]):
        """Add comment to task (stored in history)"""
        comment = action.get("comment", "")
        if comment:
            # Store as task history
            pool = await opc_db.get_pool()
            async with pool.acquire() as conn:
                await conn.execute(
                    """
                    INSERT INTO task_history (task_id, action, actor, actor_type, new_value)
                    VALUES ($1, $2, $3, $4, $5)
                """,
                    task["id"],
                    "comment",
                    "agent",
                    "agent",
                    comment,
                )

    async def _update_execution_status(self, execution_id: int, status: str):
        """Update execution status"""
        await opc_db.update_execution_status(execution_id, status)

    async def _mark_execution_completed(self, execution_id: int, result: dict[str, Any]):
        """Mark execution as completed"""
        pool = await opc_db.get_pool()
        async with pool.acquire() as conn:
            await conn.execute(
                """
                UPDATE agent_executions
                SET status = $1,
                    output_result = $2,
                    actions_proposed = $3,
                    actions_executed = $3,
                    completed_at = CURRENT_TIMESTAMP
                WHERE id = $4
            """,
                "completed",
                json.dumps(result),
                json.dumps(result.get("actions", [])),
                execution_id,
            )

    async def _mark_execution_failed(self, execution_id: int, error: str):
        """Mark execution as failed"""
        pool = await opc_db.get_pool()
        async with pool.acquire() as conn:
            await conn.execute(
                """
                UPDATE agent_executions
                SET status = $1, error_message = $2, completed_at = CURRENT_TIMESTAMP
                WHERE id = $3
            """,
                "failed",
                error,
                execution_id,
            )

    async def _mark_pending_approval(self, execution_id: int, result: dict[str, Any]):
        """Mark execution as pending approval"""
        pool = await opc_db.get_pool()
        async with pool.acquire() as conn:
            await conn.execute(
                """
                UPDATE agent_executions
                SET status = $1,
                    output_result = $2,
                    actions_proposed = $3,
                    requires_approval = TRUE
                WHERE id = $4
            """,
                "pending_approval",
                json.dumps(result),
                json.dumps(result.get("actions", [])),
                execution_id,
            )

    async def _send_telegram(self, message: str):
        """Send Telegram notification"""
        try:
            proxy_url = os.environ.get("TELEGRAM_PROXY")
            async with httpx.AsyncClient(proxy=proxy_url, timeout=10.0) as client:
                await client.post(
                    f"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendMessage",
                    data={"chat_id": TELEGRAM_CHAT_ID, "text": message, "parse_mode": "Markdown"},
                )
        except Exception as e:
            logger.error(f"Failed to send Telegram notification: {e}")

    async def _send_approval_notification(self, agent: BaseAgent, task: dict[str, Any], result: dict[str, Any]):
        """Send notification for approval request"""
        message = f"""🤖 *Agent Approval Required*

Agent: {agent.name}
Task: {task['title']}

Reasoning: {result.get('reasoning', 'No reasoning provided')}

Proposed Actions:
"""
        for i, action in enumerate(result.get("actions", []), 1):
            message += f"{i}. {action.get('type')}: {action.get('title', action.get('message', 'N/A'))}\n"

        message += "\nPlease review and approve in the OPC dashboard."

        # Send Telegram
        await self._send_telegram(message)

        # Send Email
        await email_service.send_agent_approval_email(
            agent_name=agent.name,
            task=task,
            reasoning=result.get("reasoning", "No reasoning provided"),
            actions=result.get("actions", []),
        )

    async def _send_completion_notification(self, agent: BaseAgent, task: dict[str, Any], result: dict[str, Any]):
        """Send notification for completed execution"""
        message = f"""✅ *Agent Task Completed*

Agent: {agent.name}
Task: {task['title']}

Actions Executed: {len(result.get('actions', []))}
"""
        # Send Telegram
        await self._send_telegram(message)

        # Send Email
        await email_service.send_agent_completion_email(
            agent_name=agent.name, task=task, actions_count=len(result.get("actions", []))
        )


# Global executor instance
_executor: AgentExecutor | None = None


async def get_executor() -> AgentExecutor:
    """Get or create executor instance"""
    global _executor
    if _executor is None:
        _executor = AgentExecutor()
        await _executor.initialize()
    return _executor


async def start_executor():
    """Start the executor service - initializes and starts the background task"""
    print("start_executor() called")
    try:
        executor = await get_executor()
        print(f"Executor obtained: {executor is not None}, running: {executor.running}")
        # Don't await start() - it's an infinite loop!
        # Just call it to set running=True and start the loop
        # The task will be managed by the caller
        asyncio.create_task(executor.start())
        print("Executor background task created")
    except Exception as e:
        print(f"ERROR in start_executor: {e}")
        import traceback

        traceback.print_exc()
        raise