Initial commit

server/agents/runner.py

@@ -0,0 +1,290 @@
+"""
+agents/runner.py — Agent execution and APScheduler integration (async).
+
+Owns the AsyncIOScheduler — schedules and runs all agents (cron + manual).
+Each run is tracked in the agent_runs table with token counts.
+"""
+from __future__ import annotations
+
+import asyncio
+import logging
+from datetime import datetime, timezone
+
+from apscheduler.schedulers.asyncio import AsyncIOScheduler
+from apscheduler.triggers.cron import CronTrigger
+
+from ..agent.agent import Agent, DoneEvent, ErrorEvent
+from ..config import settings
+from ..database import credential_store
+from . import tasks as agent_store
+
+logger = logging.getLogger(__name__)
+
+
+class AgentRunner:
+    def __init__(self) -> None:
+        self._agent: Agent | None = None
+        self._scheduler = AsyncIOScheduler(timezone=settings.timezone)
+        self._running: dict[str, asyncio.Task] = {}  # run_id → asyncio.Task
+
+    def init(self, agent: Agent) -> None:
+        self._agent = agent
+
+    async def start(self) -> None:
+        """Load all enabled agents with schedules into APScheduler and start it."""
+        for agent in await agent_store.list_agents():
+            if agent["enabled"] and agent["schedule"]:
+                self._add_job(agent)
+        # Daily audit log rotation at 03:00
+        self._scheduler.add_job(
+            self._rotate_audit_log,
+            trigger=CronTrigger(hour=3, minute=0, timezone=settings.timezone),
+            id="system:audit-rotation",
+            replace_existing=True,
+            misfire_grace_time=3600,
+        )
+        self._scheduler.start()
+        logger.info("[agent-runner] Scheduler started, loaded scheduled agents")
+
+    def shutdown(self) -> None:
+        if self._scheduler.running:
+            self._scheduler.shutdown(wait=False)
+        logger.info("[agent-runner] Scheduler stopped")
+
+    def _add_job(self, agent: dict) -> None:
+        try:
+            self._scheduler.add_job(
+                self._run_agent_scheduled,
+                trigger=CronTrigger.from_crontab(
+                    agent["schedule"], timezone=settings.timezone
+                ),
+                id=f"agent:{agent['id']}",
+                args=[agent["id"]],
+                replace_existing=True,
+                misfire_grace_time=300,
+            )
+            logger.info(
+                f"[agent-runner] Scheduled agent '{agent['name']}' ({agent['schedule']})"
+            )
+        except Exception as e:
+            logger.error(
+                f"[agent-runner] Failed to schedule agent '{agent['name']}': {e}"
+            )
+
+    def reschedule(self, agent: dict) -> None:
+        job_id = f"agent:{agent['id']}"
+        try:
+            self._scheduler.remove_job(job_id)
+        except Exception:
+            pass
+        if agent["enabled"] and agent["schedule"]:
+            self._add_job(agent)
+
+    def remove(self, agent_id: str) -> None:
+        try:
+            self._scheduler.remove_job(f"agent:{agent_id}")
+        except Exception:
+            pass
+
+    # ── Execution ─────────────────────────────────────────────────────────────
+
+    async def run_agent_now(self, agent_id: str, override_message: str | None = None) -> dict:
+        """UI-triggered run — bypasses schedule, returns run dict."""
+        return await self._run_agent(agent_id, ignore_rate_limit=True, override_message=override_message)
+
+    async def run_agent_and_wait(
+        self,
+        agent_id: str,
+        override_message: str,
+        session_id: str | None = None,
+        extra_tools: list | None = None,
+        force_only_extra_tools: bool = False,
+    ) -> str:
+        """Run an agent, wait for it to finish, and return the final response text."""
+        run = await self._run_agent(
+            agent_id,
+            ignore_rate_limit=True,
+            override_message=override_message,
+            session_id=session_id,
+            extra_tools=extra_tools,
+            force_only_extra_tools=force_only_extra_tools,
+        )
+        if "id" not in run:
+            logger.warning("[agent-runner] run_agent_and_wait failed for agent %s: %s", agent_id, run.get("error"))
+            return f"Could not run agent: {run.get('error', 'unknown error')}"
+        run_id = run["id"]
+        task = self._running.get(run_id)
+        if task:
+            try:
+                await asyncio.wait_for(asyncio.shield(task), timeout=300)
+            except (asyncio.TimeoutError, asyncio.CancelledError):
+                pass
+        row = await agent_store.get_run(run_id)
+        return (row.get("result") or "(no response)") if row else "(no response)"
+
+    async def _rotate_audit_log(self) -> None:
+        """Called daily by APScheduler. Purges audit entries older than the configured retention."""
+        from ..audit import audit_log
+        days_str = await credential_store.get("system:audit_retention_days")
+        days = int(days_str) if days_str else 0
+        if days <= 0:
+            return
+        deleted = await audit_log.purge(older_than_days=days)
+        logger.info("[agent-runner] Audit rotation: deleted %d entries older than %d days", deleted, days)
+
+    async def _run_agent_scheduled(self, agent_id: str) -> None:
+        """Called by APScheduler — fire and forget."""
+        await self._run_agent(agent_id, ignore_rate_limit=False)
+
+    async def _run_agent(
+        self,
+        agent_id: str,
+        ignore_rate_limit: bool = False,
+        override_message: str | None = None,
+        session_id: str | None = None,
+        extra_tools: list | None = None,
+        force_only_extra_tools: bool = False,
+    ) -> dict:
+        agent_data = await agent_store.get_agent(agent_id)
+        if not agent_data:
+            logger.warning("[agent-runner] _run_agent: agent %s not found", agent_id)
+            return {"error": "Agent not found"}
+        if not agent_data["enabled"] and not ignore_rate_limit:
+            logger.warning("[agent-runner] _run_agent: agent %s is disabled", agent_id)
+            return {"error": "Agent is disabled"}
+
+        # Kill switch
+        if await credential_store.get("system:paused") == "1":
+            logger.warning("[agent-runner] _run_agent: system is paused")
+            return {"error": "Agent is paused"}
+
+        if self._agent is None:
+            logger.warning("[agent-runner] _run_agent: agent runner not initialized")
+            return {"error": "Agent not initialized"}
+
+        # allowed_tools is JSONB, normalised to list|None in _agent_row()
+        raw = agent_data.get("allowed_tools")
+        allowed_tools: list[str] | None = raw if raw else None
+
+        # Resolve agent owner's admin status — bash is never available to non-admin owners
+        # Also block execution if the owner account has been deactivated.
+        owner_is_admin = True
+        owner_id = agent_data.get("owner_user_id")
+        if owner_id:
+            from ..users import get_user_by_id as _get_user
+            owner = await _get_user(owner_id)
+            if owner and not owner.get("is_active", True):
+                logger.warning(
+                    "[agent-runner] Skipping agent '%s' — owner account is deactivated",
+                    agent_data["name"],
+                )
+                return {"error": "Owner account is deactivated"}
+            owner_is_admin = (owner["role"] == "admin") if owner else True
+
+        if not owner_is_admin:
+            if allowed_tools is None:
+                all_names = [t.name for t in self._agent._registry.all_tools()]
+                allowed_tools = [t for t in all_names if t != "bash"]
+            else:
+                allowed_tools = [t for t in allowed_tools if t != "bash"]
+
+        # Create run record
+        run = await agent_store.create_run(agent_id)
+        run_id = run["id"]
+
+        logger.info(
+            f"[agent-runner] Running agent '{agent_data['name']}' run={run_id[:8]}"
+        )
+
+        # Per-agent max_tool_calls override (None = use system default)
+        max_tool_calls: int | None = agent_data.get("max_tool_calls") or None
+
+        async def _execute():
+            input_tokens = 0
+            output_tokens = 0
+            final_text = ""
+            try:
+                from ..agent.agent import _build_system_prompt
+                prompt_mode = agent_data.get("prompt_mode") or "combined"
+                agent_prompt = agent_data["prompt"]
+                system_override: str | None = None
+
+                if override_message:
+                    run_message = override_message
+                    if prompt_mode == "agent_only":
+                        system_override = agent_prompt
+                    elif prompt_mode == "combined":
+                        system_override = agent_prompt + "\n\n---\n\n" + await _build_system_prompt(user_id=owner_id)
+                else:
+                    run_message = agent_prompt
+                    if prompt_mode == "agent_only":
+                        system_override = agent_prompt
+                    elif prompt_mode == "combined":
+                        system_override = agent_prompt + "\n\n---\n\n" + await _build_system_prompt(user_id=owner_id)
+
+                stream = await self._agent.run(
+                    message=run_message,
+                    session_id=session_id or f"agent:{run_id}",
+                    task_id=run_id,
+                    allowed_tools=allowed_tools,
+                    model=agent_data.get("model") or None,
+                    max_tool_calls=max_tool_calls,
+                    system_override=system_override,
+                    user_id=owner_id,
+                    extra_tools=extra_tools,
+                    force_only_extra_tools=force_only_extra_tools,
+                )
+                async for event in stream:
+                    if isinstance(event, DoneEvent):
+                        final_text = event.text or "Done"
+                        input_tokens = event.usage.input_tokens
+                        output_tokens = event.usage.output_tokens
+                    elif isinstance(event, ErrorEvent):
+                        final_text = f"Error: {event.message}"
+
+                await agent_store.finish_run(
+                    run_id,
+                    status="success",
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
+                    result=final_text,
+                )
+                logger.info(
+                    f"[agent-runner] Agent '{agent_data['name']}' run={run_id[:8]} completed OK"
+                )
+            except asyncio.CancelledError:
+                await agent_store.finish_run(run_id, status="stopped")
+                logger.info(f"[agent-runner] Run {run_id[:8]} stopped")
+            except Exception as e:
+                logger.error(f"[agent-runner] Run {run_id[:8]} failed: {e}")
+                await agent_store.finish_run(run_id, status="error", error=str(e))
+            finally:
+                self._running.pop(run_id, None)
+
+        task = asyncio.create_task(_execute())
+        self._running[run_id] = task
+        return await agent_store.get_run(run_id)
+
+    def stop_run(self, run_id: str) -> bool:
+        task = self._running.get(run_id)
+        if task and not task.done():
+            task.cancel()
+            return True
+        return False
+
+    def is_running(self, run_id: str) -> bool:
+        task = self._running.get(run_id)
+        return task is not None and not task.done()
+
+    async def find_active_run(self, agent_id: str) -> str | None:
+        """Return run_id of an in-progress run for this agent, or None."""
+        for run_id, task in self._running.items():
+            if not task.done():
+                run = await agent_store.get_run(run_id)
+                if run and run["agent_id"] == agent_id:
+                    return run_id
+        return None
+
+
+# Module-level singleton
+agent_runner = AgentRunner()

server/agents/tasks.py

@@ -0,0 +1,225 @@
+"""
+agents/tasks.py — Agent and agent run CRUD operations (async).
+"""
+from __future__ import annotations
+
+import json
+import uuid
+from datetime import datetime, timezone
+from typing import Any
+
+from ..database import _rowcount, get_pool
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _agent_row(row) -> dict:
+    """Convert asyncpg Record to a plain dict, normalising JSONB fields."""
+    d = dict(row)
+    # allowed_tools: JSONB column, but SQLite-migrated rows may have stored a
+    # JSON string instead of a JSON array — asyncpg then returns a str.
+    at = d.get("allowed_tools")
+    if isinstance(at, str):
+        try:
+            d["allowed_tools"] = json.loads(at)
+        except (json.JSONDecodeError, ValueError):
+            d["allowed_tools"] = None
+    return d
+
+
+# ── Agents ────────────────────────────────────────────────────────────────────
+
+async def create_agent(
+    name: str,
+    prompt: str,
+    model: str,
+    description: str = "",
+    can_create_subagents: bool = False,
+    allowed_tools: list[str] | None = None,
+    schedule: str | None = None,
+    enabled: bool = True,
+    parent_agent_id: str | None = None,
+    created_by: str = "user",
+    max_tool_calls: int | None = None,
+    prompt_mode: str = "combined",
+    owner_user_id: str | None = None,
+) -> dict:
+    agent_id = str(uuid.uuid4())
+    now = _now()
+    pool = await get_pool()
+    await pool.execute(
+        """
+        INSERT INTO agents
+            (id, name, description, prompt, model, can_create_subagents,
+             allowed_tools, schedule, enabled, parent_agent_id, created_by,
+             created_at, updated_at, max_tool_calls, prompt_mode, owner_user_id)
+        VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16)
+        """,
+        agent_id, name, description, prompt, model,
+        can_create_subagents,
+        allowed_tools,  # JSONB — pass list directly
+        schedule, enabled,
+        parent_agent_id, created_by, now, now,
+        max_tool_calls, prompt_mode, owner_user_id,
+    )
+    return await get_agent(agent_id)
+
+
+async def list_agents(
+    include_subagents: bool = True,
+    owner_user_id: str | None = None,
+) -> list[dict]:
+    pool = await get_pool()
+    clauses: list[str] = []
+    params: list[Any] = []
+    n = 1
+
+    if not include_subagents:
+        clauses.append("parent_agent_id IS NULL")
+    if owner_user_id is not None:
+        clauses.append(f"owner_user_id = ${n}"); params.append(owner_user_id); n += 1
+
+    where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
+    rows = await pool.fetch(
+        f"""
+        SELECT a.*,
+               (SELECT started_at FROM agent_runs
+                WHERE agent_id = a.id
+                ORDER BY started_at DESC LIMIT 1) AS last_run_at
+        FROM agents a {where} ORDER BY a.created_at DESC
+        """,
+        *params,
+    )
+    return [_agent_row(r) for r in rows]
+
+
+async def get_agent(agent_id: str) -> dict | None:
+    pool = await get_pool()
+    row = await pool.fetchrow("SELECT * FROM agents WHERE id = $1", agent_id)
+    return _agent_row(row) if row else None
+
+
+async def update_agent(agent_id: str, **fields) -> dict | None:
+    if not await get_agent(agent_id):
+        return None
+    now = _now()
+    fields["updated_at"] = now
+
+    # No bool→int conversion needed — PostgreSQL BOOLEAN accepts Python bool directly
+    # No json.dumps needed — JSONB accepts Python list directly
+
+    set_parts = []
+    values: list[Any] = []
+    for i, (k, v) in enumerate(fields.items(), start=1):
+        set_parts.append(f"{k} = ${i}")
+        values.append(v)
+
+    id_param = len(fields) + 1
+    values.append(agent_id)
+
+    pool = await get_pool()
+    await pool.execute(
+        f"UPDATE agents SET {', '.join(set_parts)} WHERE id = ${id_param}", *values
+    )
+    return await get_agent(agent_id)
+
+
+async def delete_agent(agent_id: str) -> bool:
+    pool = await get_pool()
+    async with pool.acquire() as conn:
+        async with conn.transaction():
+            await conn.execute("DELETE FROM agent_runs WHERE agent_id = $1", agent_id)
+            await conn.execute(
+                "UPDATE agents SET parent_agent_id = NULL WHERE parent_agent_id = $1", agent_id
+            )
+            await conn.execute(
+                "UPDATE scheduled_tasks SET agent_id = NULL WHERE agent_id = $1", agent_id
+            )
+            status = await conn.execute("DELETE FROM agents WHERE id = $1", agent_id)
+    return _rowcount(status) > 0
+
+
+# ── Agent runs ────────────────────────────────────────────────────────────────
+
+async def create_run(agent_id: str) -> dict:
+    run_id = str(uuid.uuid4())
+    now = _now()
+    pool = await get_pool()
+    await pool.execute(
+        "INSERT INTO agent_runs (id, agent_id, started_at, status) VALUES ($1, $2, $3, 'running')",
+        run_id, agent_id, now,
+    )
+    return await get_run(run_id)
+
+
+async def finish_run(
+    run_id: str,
+    status: str,
+    input_tokens: int = 0,
+    output_tokens: int = 0,
+    result: str | None = None,
+    error: str | None = None,
+) -> dict | None:
+    now = _now()
+    pool = await get_pool()
+    await pool.execute(
+        """
+        UPDATE agent_runs
+        SET ended_at = $1, status = $2, input_tokens = $3,
+            output_tokens = $4, result = $5, error = $6
+        WHERE id = $7
+        """,
+        now, status, input_tokens, output_tokens, result, error, run_id,
+    )
+    return await get_run(run_id)
+
+
+async def get_run(run_id: str) -> dict | None:
+    pool = await get_pool()
+    row = await pool.fetchrow("SELECT * FROM agent_runs WHERE id = $1", run_id)
+    return dict(row) if row else None
+
+
+async def cleanup_stale_runs() -> int:
+    """Mark any runs still in 'running' state as 'error' (interrupted by restart)."""
+    now = _now()
+    pool = await get_pool()
+    status = await pool.execute(
+        """
+        UPDATE agent_runs
+        SET status = 'error', ended_at = $1, error = 'Interrupted by server restart'
+        WHERE status = 'running'
+        """,
+        now,
+    )
+    return _rowcount(status)
+
+
+async def list_runs(
+    agent_id: str | None = None,
+    since: str | None = None,
+    status: str | None = None,
+    limit: int = 200,
+) -> list[dict]:
+    clauses: list[str] = []
+    params: list[Any] = []
+    n = 1
+
+    if agent_id:
+        clauses.append(f"agent_id = ${n}"); params.append(agent_id); n += 1
+    if since:
+        clauses.append(f"started_at >= ${n}"); params.append(since); n += 1
+    if status:
+        clauses.append(f"status = ${n}"); params.append(status); n += 1
+
+    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
+    params.append(limit)
+
+    pool = await get_pool()
+    rows = await pool.fetch(
+        f"SELECT * FROM agent_runs {where} ORDER BY started_at DESC LIMIT ${n}",
+        *params,
+    )
+    return [dict(r) for r in rows]