chore: import upstream snapshot with attribution

application/api/user/scheduler_worker.py

@@ -0,0 +1,434 @@
+"""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,
+    }