scheduler_worker.py

"""Body of ``execute_scheduled_run`` — runs a single agent execution.

Not a DURABLE_TASK: agent runs have side effects (messages, CRM writes)
and blind auto-retry would double them. Failures after agent.gen starts
are terminal and recorded; only the pre-start load is retry-safe.
"""

from __future__ import annotations

import logging
from datetime import datetime, timezone
from typing import Any, Dict, Optional

from celery.exceptions import SoftTimeLimitExceeded
from sqlalchemy import text as sql_text

from application.agents.headless_runner import run_agent_headless
from application.core.settings import settings
from application.events.publisher import publish_user_event
from application.storage.db.base_repository import row_to_dict
from application.storage.db.engine import get_engine
from application.storage.db.repositories.conversations import (
    ConversationsRepository,
)
from application.storage.db.repositories.schedule_runs import (
    ScheduleRunsRepository,
)
from application.storage.db.repositories.schedules import SchedulesRepository
from application.storage.db.repositories.token_usage import TokenUsageRepository

logger = logging.getLogger(__name__)


# Cap output verbatim in the run log; beyond the cap we keep the head and stamp output_truncated.
_OUTPUT_CAP_CHARS = 24_000


def _agent_config_for_schedule(schedule: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """Resolve the agent row (agent-bound) or build an ephemeral classic config.

    For agentless schedules (``agent_id IS NULL``), the worker constructs an
    in-memory agent shape carrying just enough fields for ``run_agent_headless``:
    classic agent type, system-default retriever/chunks/prompt, no source, and
    the optional ``model_id`` override. The runtime toolset is rebuilt by
    ``ToolExecutor`` at fire time (current ``user_tools`` + non-disabled,
    non-headless-excluded defaults), so a snapshot here would be dead code.
    """
    if schedule.get("agent_id"):
        engine = get_engine()
        with engine.connect() as conn:
            row = conn.execute(
                sql_text("SELECT * FROM agents WHERE id = CAST(:id AS uuid)"),
                {"id": str(schedule["agent_id"])},
            ).fetchone()
        return row_to_dict(row) if row is not None else None
    return _ephemeral_agent_for_agentless(schedule)


def _ephemeral_agent_for_agentless(
    schedule: Dict[str, Any],
) -> Optional[Dict[str, Any]]:
    """Build an agent-shaped config for a schedule with no parent agent."""
    # ``agent_config["tools"]`` is intentionally omitted: ``run_agent_headless``
    # never reads it. The runtime toolset is rebuilt by
    # ``ToolExecutor._get_user_tools(owner)`` at fire time — same dereference
    # the agent-bound path uses, so a tool added/disabled after creation is
    # reflected. Headless mode there filters chat-only tools (``scheduler``).
    user_id = schedule.get("user_id")
    if not user_id:
        return None
    return {
        "id": None,
        "user_id": user_id,
        "agent_type": "classic",
        "retriever": "classic",
        "chunks": 2,
        "prompt_id": "default",
        "source_id": None,
        "default_model_id": schedule.get("model_id") or "",
    }


def _load_chat_history(schedule: Dict[str, Any]) -> list:
    """Originating conversation history (one-time only; recurring has none)."""
    origin = schedule.get("origin_conversation_id")
    if not origin or schedule.get("trigger_type") != "once":
        return []
    user_id = schedule.get("user_id")
    if not user_id:
        return []
    try:
        engine = get_engine()
        with engine.connect() as conn:
            conv = ConversationsRepository(conn).get_any(str(origin), user_id)
            if conv is None:
                return []
            messages = ConversationsRepository(conn).get_messages(str(conv["id"]))
    except Exception:
        logger.exception("scheduler: failed loading chat history")
        return []
    history: list = []
    for msg in messages:
        if msg.get("prompt") and msg.get("response"):
            history.append({
                "prompt": msg["prompt"],
                "response": msg["response"],
            })
    return history


def _publish_run_event(
    event_type: str, run: Dict[str, Any], schedule: Dict[str, Any], **extra: Any
) -> None:
    """Best-effort SSE publish for a scheduler run state transition."""
    user_id = run.get("user_id") or schedule.get("user_id")
    if not user_id:
        return
    agent_id_raw = schedule.get("agent_id")
    payload = {
        "run_id": str(run["id"]),
        "schedule_id": str(schedule["id"]),
        "agent_id": str(agent_id_raw) if agent_id_raw else None,
        "trigger_type": schedule.get("trigger_type"),
        "status": run.get("status"),
        **extra,
    }
    try:
        publish_user_event(
            user_id,
            event_type,
            payload,
            scope={"kind": "schedule", "id": str(schedule["id"])},
        )
    except Exception:
        logger.exception(
            "scheduler: SSE publish failed event=%s run=%s",
            event_type, run.get("id"),
        )


def _publish_message_appended(
    user_id: str,
    conversation_id: str,
    message: Dict[str, Any],
    schedule_id: str,
    run_id: str,
) -> None:
    """SSE message-appended event for a one-time run's chat turn."""
    try:
        publish_user_event(
            user_id,
            "schedule.message.appended",
            {
                "conversation_id": str(conversation_id),
                "message_id": str(message["id"]),
                "schedule_id": str(schedule_id),
                "run_id": str(run_id),
                "position": int(message.get("position", 0)),
            },
            scope={"kind": "conversation", "id": str(conversation_id)},
        )
    except Exception:
        logger.exception(
            "scheduler: message.appended publish failed run=%s", run_id,
        )


def _append_one_time_turn(
    schedule: Dict[str, Any],
    run: Dict[str, Any],
    outcome: Dict[str, Any],
) -> Optional[Dict[str, Any]]:
    """Insert an assistant turn in the originating conversation (once only)."""
    origin = schedule.get("origin_conversation_id")
    if not origin:
        return None
    engine = get_engine()
    user_id = schedule.get("user_id")
    metadata = {
        "scheduled": True,
        "schedule_id": str(schedule["id"]),
        "run_id": str(run["id"]),
        "scheduled_run_at": (
            run.get("scheduled_for")
            if isinstance(run.get("scheduled_for"), str)
            else None
        ),
    }
    with engine.begin() as conn:
        conv = ConversationsRepository(conn).get_any(str(origin), user_id)
        if conv is None:
            return None
        message = ConversationsRepository(conn).append_message(
            str(conv["id"]),
            {
                "prompt": schedule.get("instruction") or "",
                "response": outcome.get("answer") or "",
                "thought": outcome.get("thought") or "",
                "sources": outcome.get("sources") or [],
                "tool_calls": outcome.get("tool_calls") or [],
                "model_id": outcome.get("model_id"),
                "metadata": metadata,
            },
        )
    return message


def execute_scheduled_run_body(run_id: str, celery_task_id: Optional[str]) -> Dict[str, Any]:
    """Execute one scheduled run by id; returns a result dict for tracing."""
    if not settings.POSTGRES_URI:
        return {"status": "skipped", "reason": "POSTGRES_URI not set"}

    engine = get_engine()

    with engine.connect() as conn:
        run = ScheduleRunsRepository(conn).get_internal(run_id)
        if run is None:
            return {"status": "skipped", "reason": "run not found"}
        schedule = SchedulesRepository(conn).get_internal(str(run["schedule_id"]))
    if schedule is None:
        return {"status": "skipped", "reason": "schedule not found"}

    # Refuse non-runnable terminal states; manual run-now bypasses.
    if run.get("status") != "pending":
        return {"status": "skipped", "reason": f"run status={run.get('status')}"}
    if schedule.get("status") in {"cancelled", "completed"} and run.get(
        "trigger_source"
    ) != "manual":
        with engine.begin() as conn:
            ScheduleRunsRepository(conn).update(
                run_id,
                {
                    "status": "skipped",
                    "finished_at": datetime.now(timezone.utc),
                    "error_type": "internal",
                    "error": "schedule no longer active",
                },
            )
        return {"status": "skipped", "reason": "schedule terminal"}

    agent_config = _agent_config_for_schedule(schedule)
    if agent_config is None:
        with engine.begin() as conn:
            updated = ScheduleRunsRepository(conn).update(
                run_id,
                {
                    "status": "failed",
                    "finished_at": datetime.now(timezone.utc),
                    "error_type": "internal",
                    "error": "agent missing",
                },
            )
            SchedulesRepository(conn).bump_failure_count(str(schedule["id"]))
        _publish_run_event("schedule.run.failed", updated or run, schedule,
                           error="agent missing")
        return {"status": "failed", "reason": "agent missing"}

    with engine.begin() as conn:
        if not ScheduleRunsRepository(conn).mark_running(run_id, celery_task_id):
            return {"status": "skipped", "reason": "lost race to mark_running"}

    started = datetime.now(timezone.utc)
    instruction = schedule.get("instruction") or ""
    allowlist = schedule.get("tool_allowlist") or []
    chat_history = _load_chat_history(schedule)
    outcome: Dict[str, Any]
    error_type: Optional[str] = None
    error_text: Optional[str] = None
    timed_out = False
    try:
        outcome = run_agent_headless(
            agent_config,
            instruction,
            tool_allowlist=allowlist,
            model_id_override=schedule.get("model_id"),
            endpoint="schedule",
            conversation_id=schedule.get("origin_conversation_id"),
            chat_history=chat_history,
        )
    except SoftTimeLimitExceeded:
        timed_out = True
        outcome = {"answer": "", "tool_calls": [], "sources": [], "thought": ""}
        error_type = "timeout"
        error_text = "run exceeded soft time limit"
    except Exception as exc:
        outcome = {"answer": "", "tool_calls": [], "sources": [], "thought": ""}
        error_type = "agent_error"
        error_text = str(exc)
        logger.exception("scheduler: agent run failed run=%s", run_id)

    finished = datetime.now(timezone.utc)

    # Headless denial with no usable output → tool_not_allowed.
    if (
        error_type is None
        and (outcome.get("denied") or [])
        and not (outcome.get("answer") or "").strip()
    ):
        error_type = "tool_not_allowed"
        error_text = "headless allowlist blocked required tool"

    prompt_tokens = int(outcome.get("prompt_tokens", 0) or 0)
    generated_tokens = int(outcome.get("generated_tokens", 0) or 0)
    used_tokens = prompt_tokens + generated_tokens
    if (
        schedule.get("token_budget") is not None
        and int(schedule["token_budget"]) > 0
        and used_tokens > int(schedule["token_budget"])
    ):
        error_type = "budget_exceeded"
        error_text = (
            f"used {used_tokens} tokens exceeds budget "
            f"{schedule['token_budget']}"
        )

    answer = outcome.get("answer") or ""
    truncated = False
    if len(answer) > _OUTPUT_CAP_CHARS:
        answer = answer[:_OUTPUT_CAP_CHARS]
        truncated = True

    new_status = (
        "timeout" if timed_out else ("failed" if error_type else "success")
    )

    with engine.begin() as conn:
        update_fields: Dict[str, Any] = {
            "status": new_status,
            "started_at": started,
            "finished_at": finished,
            "output": answer or None,
            "output_truncated": truncated,
            "prompt_tokens": prompt_tokens,
            "generated_tokens": generated_tokens,
        }
        if error_type:
            update_fields["error_type"] = error_type
            update_fields["error"] = error_text
        updated_run = ScheduleRunsRepository(conn).update(run_id, update_fields)
        if used_tokens > 0:
            agent_id_raw = schedule.get("agent_id")
            try:
                TokenUsageRepository(conn).insert(
                    user_id=schedule.get("user_id"),
                    api_key=None,
                    prompt_tokens=prompt_tokens,
                    generated_tokens=generated_tokens,
                    timestamp=finished,
                    agent_id=str(agent_id_raw) if agent_id_raw else None,
                    source="schedule",
                    request_id=str(run_id),
                    model_id=outcome.get("model_id"),
                )
            except Exception:
                logger.exception(
                    "scheduler: token_usage insert failed run=%s", run_id,
                )
        schedules_repo = SchedulesRepository(conn)
        autopaused = False
        if new_status == "success":
            schedules_repo.reset_failure_count(str(schedule["id"]))
        elif new_status in ("failed", "timeout"):
            count = schedules_repo.bump_failure_count(str(schedule["id"]))
            if (
                settings.SCHEDULE_AUTOPAUSE_FAILURES > 0
                and count >= settings.SCHEDULE_AUTOPAUSE_FAILURES
                and schedule.get("trigger_type") == "recurring"
            ):
                autopaused = schedules_repo.autopause(str(schedule["id"]))
        # Once: terminal-flip on cron-fired runs only; manual runs on a
        # still-active once-schedule leave the future cadence intact.
        if (
            schedule.get("trigger_type") == "once"
            and run.get("trigger_source") != "manual"
            and schedule.get("status") == "active"
        ):
            schedules_repo.update_internal(
                str(schedule["id"]),
                {"status": "completed", "next_run_at": None},
            )

    appended: Optional[Dict[str, Any]] = None
    if (
        schedule.get("trigger_type") == "once"
        and new_status == "success"
        and schedule.get("origin_conversation_id")
    ):
        try:
            appended = _append_one_time_turn(schedule, updated_run or run, outcome)
        except Exception:
            logger.exception(
                "scheduler: append turn failed run=%s", run_id,
            )
        if appended is not None:
            with engine.begin() as conn:
                ScheduleRunsRepository(conn).update(
                    run_id,
                    {
                        "conversation_id": str(appended["conversation_id"]),
                        "message_id": str(appended["id"]),
                    },
                )
            _publish_message_appended(
                schedule.get("user_id"),
                str(appended["conversation_id"]),
                appended,
                str(schedule["id"]),
                run_id,
            )

    if new_status == "success":
        _publish_run_event("schedule.run.completed", updated_run or run, schedule)
    else:
        _publish_run_event(
            "schedule.run.failed",
            updated_run or run,
            schedule,
            error_type=error_type,
            error=error_text,
        )

    if autopaused:
        _publish_run_event(
            "schedule.autopaused",
            updated_run or run,
            schedule,
            consecutive_failure_count=settings.SCHEDULE_AUTOPAUSE_FAILURES,
        )

    return {
        "status": new_status,
        "run_id": run_id,
        "error_type": error_type,
    }