apps/api/app/services/runner.py

"""Run orchestration: start a run, pause at the human gate, resume on decision.

The graph is invoked twice per brief:

1. ``ainvoke(initial_state)`` runs supervisor → workers → writer → verify → gate,
   where :func:`langgraph.types.interrupt` suspends it and checkpoints state.
2. ``ainvoke(Command(resume=decision))`` continues from that checkpoint through
   delivery.

Between the two the process could restart entirely — that is the point of the
checkpointer, and why approval is a separate authenticated HTTP request rather
than a callback held in memory.
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import re
import time
import uuid
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, cast

from langgraph.graph.state import CompiledStateGraph
from langgraph.types import Command

from app.core.budget import BudgetGuard
from app.core.claude import build_engine
from app.core.events import EventBus, EventKind, get_event_bus
from app.core.settings import Settings, get_settings
from app.core.tracing import Tracer, get_tracer
from app.graph.build import build_graph
from app.graph.state import RunState, initial_state, is_partial
from app.mcp_server.client import McpToolClient, ToolCallRecord
from app.models.brief import Brief
from app.models.run import HumanDecision, RunMode, RunStatus, TokenSpend
from app.services.render import render_markdown
from app.services.repository import Repository, get_repository
from app.services.serde import build_checkpoint_serde

logger = logging.getLogger(__name__)

#: Symbol injected by `demo_fault` mode. Verified to return no provider data.
DEMO_FAULT_TICKER = "ZZZZQQQQ"


class RunNotFoundError(RuntimeError):
    """No such run, or the run is not awaiting a decision."""


class RunNotResumableError(RuntimeError):
    """The run is awaiting a decision but its checkpoint is gone.

    Raised rather than silently re-invoking the graph from scratch: a
    ``Command(resume=...)`` against an unknown thread does not fail loudly, it
    starts a brand-new execution whose state never reaches the gate.
    """


#: Reason recorded against runs abandoned by a previous process.
ORPHAN_REASON = (
    "abandoned by a server restart — no checkpoint survived, so this run could "
    "not be resumed and was closed to release its watchlist"
)


def watchlist_key_for(tickers: list[str]) -> str:
    """Canonical key for the active-run constraint."""
    return ",".join(sorted({t.strip().upper() for t in tickers if t.strip()}))


def sqlite_checkpoint_path(database_url: str) -> Path | None:
    """Where to checkpoint for a SQLite archive URL, or ``None`` if not on disk.

    An in-memory database (``sqlite://`` or ``sqlite:///:memory:``) has nothing
    to be durable *to*, so it keeps the in-memory saver and reports itself as
    volatile rather than pretending otherwise.
    """
    match = re.match(r"^sqlite(?:\+\w+)?://", database_url)
    if match is None:
        return None
    # SQLAlchemy spells SQLite as `sqlite://<host>/<path>` with an empty host,
    # so exactly one slash separates scheme from path: `sqlite:///rel.db` is
    # relative and `sqlite:////abs.db` is absolute. Stripping more than one
    # would silently turn every absolute path into a relative one.
    location = database_url[match.end() :].removeprefix("/")
    if not location or location == ":memory:":
        return None
    path = Path(location)
    return path.with_name(f"{path.stem}.checkpoints{path.suffix or '.db'}")


def _daemonise(connection: Any) -> None:
    """Let the interpreter exit even if this connection is never closed.

    aiosqlite serves every statement from a worker thread it creates *without*
    ``daemon=True``, so an unclosed connection does not merely leak — it wedges
    interpreter shutdown indefinitely, waiting on a thread that only
    ``close()`` can stop. :meth:`RunService.shutdown` does close it, but a
    missed shutdown should cost a leaked handle, not a process that never dies.

    Losing an in-flight write to an abrupt exit is acceptable here: every write
    that mattered was awaited before we reported success, and WAL keeps the file
    consistent regardless. Reaching for a private attribute is deliberate and
    guarded — if a future aiosqlite reshapes it, this quietly does nothing.
    """
    thread = getattr(connection, "_thread", None)
    if thread is not None and hasattr(thread, "daemon"):
        with contextlib.suppress(RuntimeError):
            thread.daemon = True


class RunService:
    """Owns the compiled graph, the checkpointer and the background run tasks."""

    def __init__(
        self,
        settings: Settings | None = None,
        repository: Repository | None = None,
        bus: EventBus | None = None,
        tracer: Tracer | None = None,
    ) -> None:
        self.settings = settings or get_settings()
        self.repository = repository or get_repository()
        self.bus = bus or get_event_bus()
        self.tracer = tracer or get_tracer()
        #: True only when checkpoints outlive the process (Postgres or SQLite).
        self._durable_checkpoints = False
        #: Guards the one-time async upgrade performed by `ensure_ready`.
        self._checkpointer_ready = False
        #: Held so shutdown can close it; only set on a durable path.
        self._saver_connection: Any | None = None
        self._checkpointer = self._build_checkpointer()
        self.graph: CompiledStateGraph[Any, Any, Any, Any] = build_graph(self._checkpointer)
        self._tasks: dict[str, asyncio.Task[None]] = {}
        self._started_at: dict[str, float] = {}

    @property
    def durable_checkpoints(self) -> bool:
        """Whether a run paused at the gate would survive a restart.

        Surfaced rather than inferred from the database URL: the upgrade can
        fail (an unreachable Postgres, a read-only disk) and degrade to the
        in-memory saver, and the console must report what is true, not what was
        configured.
        """
        return self._durable_checkpoints

    # ------------------------------------------------------------ checkpointer ---
    def _build_checkpointer(self) -> Any:
        """The in-process saver the service always starts with.

        Postgres cannot be set up here: the async saver needs an awaited
        connection and an awaited ``setup()``, and ``__init__`` cannot await.
        :meth:`ensure_ready` performs that upgrade during application startup.

        Every saver is given a strict serializer allowlisted to AlphaBrief's own
        models (see :mod:`app.services.serde`), so checkpoint bytes can never
        rehydrate an arbitrary type.
        """
        serde = build_checkpoint_serde()
        from langgraph.checkpoint.memory import InMemorySaver

        self._durable_checkpoints = False
        logger.info("LangGraph checkpointer: InMemorySaver (paused runs end at shutdown)")
        return InMemorySaver(serde=serde)

    async def ensure_ready(self) -> None:
        """Upgrade to a durable checkpointer for the configured database.

        Called once during application startup, before reconciliation.

        Both supported databases get a durable saver, because the human gate is
        the product: a run paused for approval that cannot survive a restart is
        a run the reviewer can lose by closing a laptop lid. Postgres is used
        when a Postgres URL is configured; otherwise the default SQLite file is
        checkpointed too, so a clean clone with no configuration at all still
        has a gate that outlives the process.

        The **async** saver is mandatory here, not a preference: the graph is
        driven by ``ainvoke``, and langgraph's synchronous ``PostgresSaver``
        inherits ``aget_tuple`` from the base class, where it raises
        ``NotImplementedError``. Pairing it with an async graph fails every run
        the moment the loop opens. ``SqliteSaver`` has the identical defect.

        Falls back to the in-memory saver on any connection problem: a database
        outage should cost durability, not the ability to produce briefs.
        """
        if self._checkpointer_ready:
            return
        self._checkpointer_ready = True

        url = self.settings.database_url
        if url.startswith(("postgres://", "postgresql://", "postgresql+")):
            await self._upgrade_to_postgres(url)
            return
        path = sqlite_checkpoint_path(url)
        if path is not None:
            await self._upgrade_to_sqlite(path)

    async def _upgrade_to_postgres(self, url: str) -> None:
        try:
            from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
            from psycopg import AsyncConnection
            from psycopg.rows import dict_row

            from app.services.repository import normalise_database_url

            dsn = normalise_database_url(url).replace("postgresql+psycopg://", "postgresql://")
            # prepare_threshold=0 because the pooled endpoint multiplexes
            # sessions; server-side prepared statements do not survive that.
            connection = await AsyncConnection.connect(
                dsn, autocommit=True, prepare_threshold=0, row_factory=dict_row
            )
            saver = AsyncPostgresSaver(connection, serde=build_checkpoint_serde())
            await saver.setup()
        except Exception:
            logger.warning(
                "Postgres checkpointer unavailable; staying on the in-memory saver",
                exc_info=True,
            )
            return

        self._adopt(saver, connection, "AsyncPostgresSaver")

    async def _upgrade_to_sqlite(self, path: Path) -> None:
        """Checkpoint into a file beside the archive database.

        A *separate* file, not the archive itself: SQLAlchemy drives the archive
        synchronously while the saver drives this one through aiosqlite, and
        pointing two connections with different drivers at one file trades a
        clear separation for `database is locked`. WAL is enabled so a reader
        never blocks the writer.
        """
        try:
            import aiosqlite
            from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver

            path.parent.mkdir(parents=True, exist_ok=True)
            connection = await aiosqlite.connect(str(path))
            _daemonise(connection)
            await connection.execute("PRAGMA journal_mode=WAL")
            saver = AsyncSqliteSaver(connection, serde=build_checkpoint_serde())
            await saver.setup()
        except Exception:
            logger.warning(
                "SQLite checkpointer unavailable; staying on the in-memory saver",
                exc_info=True,
            )
            return

        self._adopt(saver, connection, f"AsyncSqliteSaver at {path.name}")

    def _adopt(self, saver: Any, connection: Any, label: str) -> None:
        """Swap in a durable saver and rebuild the graph around it."""
        self._saver_connection = connection
        self._checkpointer = saver
        self.graph = build_graph(saver)
        self._durable_checkpoints = True
        logger.info("LangGraph checkpointer: %s (durable)", label)

    # ------------------------------------------------------------ reconcile ---
    async def _has_checkpoint(self, run_id: str) -> bool:
        """Whether the checkpointer still holds a resumable thread for this run.

        Uses the async interface because that is the one the graph itself uses —
        checking through a path the saver does not actually implement would make
        every run look unresumable.
        """
        config = {"configurable": {"thread_id": run_id}}
        try:
            return await self._checkpointer.aget_tuple(config) is not None
        except Exception:
            logger.warning("checkpoint lookup failed for %s", run_id, exc_info=True)
            return False

    async def reconcile_orphans(self) -> int:
        """Close runs abandoned by a previous process, releasing their watchlists.

        Every active status occupies its watchlist's single active-run slot, and
        no code path moves an abandoned row out of that set — so without this a
        crash mid-run makes the watchlist permanently un-runnable. A run is an
        orphan when this process is not executing it and it cannot be resumed:

        * ``QUEUED``/``RUNNING`` — nothing resumes mid-graph work, so an inherited
          row is always an orphan regardless of checkpoint durability.
        * ``AWAITING_APPROVAL``/``HUMAN_REVIEW`` — genuinely resumable, but only
          while a checkpoint survives. With Postgres it does and the run is kept.
        """
        try:
            active = await self.repository.active_runs()
        except Exception:
            logger.warning("orphan reconciliation skipped: could not list runs", exc_info=True)
            return 0

        closed = 0
        for record in active:
            run_id = record["id"]
            if run_id in self._tasks:
                continue
            status = record["status"]
            resumable = status in (RunStatus.AWAITING_APPROVAL, RunStatus.HUMAN_REVIEW)
            if resumable and self._durable_checkpoints and await self._has_checkpoint(run_id):
                logger.info("run %s kept: durable checkpoint intact at the gate", run_id)
                continue
            await self.repository.finish_run(
                run_id,
                status=RunStatus.FAILED,
                abort_reason=ORPHAN_REASON,
            )
            closed += 1
            logger.warning("run %s closed as orphaned (was %s)", run_id, status)

        if closed:
            logger.warning("reconciled %d orphaned run(s) from a previous process", closed)
        return closed

    # -------------------------------------------------------------------- api ---
    async def start_run(
        self,
        *,
        tickers: list[str],
        mode: RunMode = RunMode.STANDARD,
        trigger: str = "ui",
    ) -> str:
        """Register a run and kick off graph execution in the background."""
        cleaned = [t.strip().upper() for t in tickers if t.strip()]
        if mode == RunMode.DEMO_FAULT and DEMO_FAULT_TICKER not in cleaned:
            cleaned.append(DEMO_FAULT_TICKER)

        run_id = f"run_{uuid.uuid4().hex[:20]}"
        key = watchlist_key_for(cleaned)

        await self.repository.create_run(
            run_id=run_id,
            watchlist_key=key,
            tickers=cleaned,
            mode=str(mode),
            engine=self.settings.resolved_engine,
            trigger=trigger,
        )

        await self.bus.publish(
            run_id,
            EventKind.RUN_STARTED,
            f"Run started for {', '.join(cleaned)}",
            {
                "tickers": cleaned,
                "mode": str(mode),
                "engine": self.settings.resolved_engine,
                "trigger": trigger,
            },
        )

        self._started_at[run_id] = time.perf_counter()
        task = asyncio.create_task(self._execute(run_id, cleaned, key, mode))
        self._tasks[run_id] = task
        task.add_done_callback(lambda _t: self._tasks.pop(run_id, None))
        return run_id

    async def submit_decision(self, run_id: str, decision: HumanDecision) -> dict[str, Any]:
        """Resume a paused graph with the reviewer's decision."""
        record = await self.repository.get_run(run_id)
        if record is None:
            raise RunNotFoundError(f"run '{run_id}' not found")
        if record["status"] not in (RunStatus.AWAITING_APPROVAL, RunStatus.HUMAN_REVIEW):
            raise RunNotFoundError(f"run '{run_id}' is {record['status']}, not awaiting a decision")

        # A resume against a missing thread does not raise — LangGraph would start
        # a fresh execution that never reaches the gate, and the run would land in
        # FAILED with no explanation. Refuse explicitly instead.
        if not await self._has_checkpoint(run_id):
            await self.repository.finish_run(
                run_id,
                status=RunStatus.FAILED,
                abort_reason=ORPHAN_REASON,
            )
            raise RunNotResumableError(
                f"run '{run_id}' can no longer be approved: the server restarted and its "
                "checkpoint was lost. The brief is still in the archive. Configure "
                "DATABASE_URL with Postgres to make paused runs survive a restart."
            )

        await self.repository.save_approval(
            run_id=run_id,
            action=decision.action,
            reviewer=decision.reviewer,
            note=decision.note,
        )

        budget = BudgetGuard(self.settings)
        async with McpToolClient(settings=self.settings) as mcp:
            ctx = self._context(run_id, mcp, budget)
            final = await self.graph.ainvoke(
                Command(resume=decision.model_dump()),
                ctx.to_config(),
                context=ctx,
            )

        await self._finalise(run_id, final)
        return {"run_id": run_id, "status": final.get("status"), "action": decision.action}

    async def get_run(self, run_id: str) -> dict[str, Any] | None:
        return await self.repository.get_run(run_id)

    async def list_runs(self, limit: int = 50) -> list[dict[str, Any]]:
        return await self.repository.list_runs(limit)

    async def gate_payload(self, run_id: str) -> dict[str, Any] | None:
        """The approval-gate view for a paused run."""
        archived = await self.repository.latest_brief(run_id)
        if archived is None:
            return None
        record = await self.repository.get_run(run_id)
        return {
            "run_id": run_id,
            "status": record["status"] if record else None,
            "verified": archived["verified"],
            "requires_review": not archived["verified"],
            "brief": archived["brief"],
            "verification": archived["verification"],
            "markdown": archived["markdown"],
        }

    async def shutdown(self) -> None:
        for task in list(self._tasks.values()):
            task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await task
        self.tracer.flush()
        if self._saver_connection is not None:
            with contextlib.suppress(Exception):
                await self._saver_connection.close()
            self._saver_connection = None

    # --------------------------------------------------------------- internal ---
    def _context(self, run_id: str, mcp: McpToolClient, budget: BudgetGuard) -> Any:
        from app.graph.context import RunContext

        ctx = RunContext(
            run_id=run_id,
            settings=self.settings,
            engine=build_engine(self.settings),
            mcp=mcp,
            bus=self.bus,
            tracer=self.tracer,
            budget=budget,
        )

        async def _emit_tool_call(record: ToolCallRecord) -> None:
            await ctx.emit(
                EventKind.MCP_TOOL_CALL,
                f"{record.tool} · {record.duration_ms:.0f}ms",
                record.to_dict(),
            )

        mcp.emitter = _emit_tool_call
        return ctx

    async def _execute(
        self, run_id: str, tickers: list[str], watchlist_key: str, mode: RunMode
    ) -> None:
        """Background task: run the graph up to the human gate."""
        budget = BudgetGuard(self.settings)
        session_date = datetime.now(UTC).date().isoformat()
        state = initial_state(
            run_id=run_id,
            watchlist_key=watchlist_key,
            tickers=tickers,
            session_date=session_date,
            mode=mode,
            days=self.settings.price_history_days,
            news_limit=self.settings.news_limit,
            max_regenerations=self.settings.max_regenerations,
        )

        try:
            async with McpToolClient(settings=self.settings) as mcp:
                ctx = self._context(run_id, mcp, budget)
                result = await asyncio.wait_for(
                    self.graph.ainvoke(state, ctx.to_config(), context=ctx),
                    timeout=self.settings.run_timeout_seconds,
                )
        except TimeoutError:
            await self._fail(run_id, "run exceeded its wall-clock timeout")
            return
        except Exception as exc:
            logger.exception("run %s crashed", run_id)
            await self._fail(run_id, f"{type(exc).__name__}: {exc}")
            return

        await self._after_first_pass(run_id, result)

    async def _after_first_pass(self, run_id: str, result: dict[str, Any]) -> None:
        """Persist the brief and mark the run as awaiting a human decision."""
        status = result.get("status")
        interrupted = "__interrupt__" in result

        if not interrupted:
            # Aborted before the gate (budget or iteration cap), or failed.
            terminal = (
                status
                if status
                in (
                    RunStatus.BUDGET_ABORT,
                    RunStatus.ITERATION_ABORT,
                )
                else RunStatus.FAILED
            )
            await self._persist_metrics(run_id, result)
            await self.repository.finish_run(
                run_id,
                str(terminal),
                abort_reason=result.get("abort_reason"),
            )
            await self.bus.publish(
                run_id,
                EventKind.RUN_FAILED,
                f"Run ended with status {terminal}",
                {"status": str(terminal), "reason": result.get("abort_reason")},
            )
            await self._archive_events(run_id)
            return

        brief = result.get("brief")
        report = result.get("verification") or {}
        verified = bool(report.get("ok"))
        gate_status = RunStatus.AWAITING_APPROVAL if verified else RunStatus.HUMAN_REVIEW

        if isinstance(brief, Brief):
            await self.repository.save_brief(
                run_id=run_id,
                generated_for=brief.generated_for,
                headline=brief.headline,
                partial=brief.partial,
                verified=verified,
                claims_total=int(report.get("checked_claims", 0)),
                claims_matched=int(report.get("matched", 0)),
                brief_json=brief.model_dump(),
                verification_json=report,
                markdown=render_markdown(brief),
            )

        await self._persist_metrics(run_id, result, status=str(gate_status), verified=verified)
        await self.bus.publish(
            run_id,
            EventKind.GATE_AWAITING,
            "Awaiting human approval"
            if verified
            else "Verification failed — flagged for human review",
            {
                "verified": verified,
                "status": str(gate_status),
                "claims_matched": report.get("matched"),
                "claims_total": report.get("checked_claims"),
                "partial": bool(brief.partial) if isinstance(brief, Brief) else True,
            },
        )

    async def _finalise(self, run_id: str, result: dict[str, Any]) -> None:
        """Persist terminal state after the gate resumes."""
        status = str(result.get("status") or RunStatus.FAILED)
        brief = result.get("brief")
        report = result.get("verification") or {}

        if isinstance(brief, Brief):
            await self.repository.save_brief(
                run_id=run_id,
                generated_for=brief.generated_for,
                headline=brief.headline,
                partial=brief.partial,
                verified=bool(report.get("ok")),
                claims_total=int(report.get("checked_claims", 0)),
                claims_matched=int(report.get("matched", 0)),
                brief_json=brief.model_dump(),
                verification_json=report,
                markdown=render_markdown(brief),
            )

        # record_latency=False: the reported latency is the machine time to
        # produce the brief, measured when it reached the gate. Everything after
        # that point is a human reading it.
        await self._persist_metrics(run_id, result, status=status, record_latency=False)
        await self.repository.finish_run(run_id, status)
        await self.bus.publish(
            run_id,
            EventKind.RUN_COMPLETED,
            f"Run {status.lower()}",
            {
                "status": status,
                "delivery": result.get("delivery"),
                "cost_usd": _spend(result).cost_usd,
            },
        )
        await self._archive_events(run_id)

    async def _elapsed_ms(self, run_id: str) -> float:
        """Milliseconds of machine work since the run was triggered.

        Prefers the in-process monotonic clock, which is immune to wall-clock
        adjustments. Falls back to the persisted `created_at`, because
        `perf_counter` is process-relative and therefore meaningless after a
        restart — and now that paused runs legitimately survive restarts, that
        fallback is a normal path rather than an edge case.
        """
        started = self._started_at.get(run_id)
        if started is not None:
            return (time.perf_counter() - started) * 1000
        record = await self.repository.get_run(run_id)
        created_raw = record.get("created_at") if record else None
        if not isinstance(created_raw, str):
            return 0.0
        try:
            created = datetime.fromisoformat(created_raw)
        except ValueError:
            return 0.0
        if created.tzinfo is None:
            created = created.replace(tzinfo=UTC)
        return max(0.0, (datetime.now(UTC) - created).total_seconds() * 1000)

    async def _persist_metrics(
        self,
        run_id: str,
        result: dict[str, Any],
        status: str | None = None,
        verified: bool | None = None,
        *,
        record_latency: bool = True,
    ) -> None:
        """Write run metrics.

        `record_latency=False` preserves the measurement taken when the brief
        reached the gate. Recomputing after approval would fold the reviewer's
        deliberation time into a figure that is reported as the system's
        latency — a five-minute coffee break would read as a five-minute brief.
        """
        spend = _spend(result)
        fields: dict[str, Any] = {
            "iterations": int(result.get("iterations", 0)),
            "model_calls": spend.calls,
            "input_tokens": spend.input_tokens,
            "output_tokens": spend.output_tokens,
            "cost_usd": spend.cost_usd,
            "tool_calls": len(result.get("tool_calls", [])),
            "partial": is_partial(_as_state(result)),
            "error_count": len(result.get("errors", [])),
        }
        if record_latency:
            fields["latency_ms"] = await self._elapsed_ms(run_id)
        if status is not None:
            fields["status"] = status
        if verified is not None:
            fields["verified"] = verified
        await self.repository.update_run(run_id, **fields)

    async def _fail(self, run_id: str, reason: str) -> None:
        # A crashed run still consumed real time; leaving latency at 0 would
        # understate the cost of a failure in the archive.
        await self.repository.finish_run(
            run_id,
            str(RunStatus.FAILED),
            abort_reason=reason,
            latency_ms=await self._elapsed_ms(run_id),
        )
        await self.bus.publish(run_id, EventKind.RUN_FAILED, reason, {"reason": reason})
        await self._archive_events(run_id)

    async def _archive_events(self, run_id: str) -> None:
        try:
            events = await self.bus.history(run_id)
            await self.repository.save_events(events)
        except Exception:
            logger.warning("failed to archive events for %s", run_id, exc_info=True)


def _spend(result: dict[str, Any]) -> TokenSpend:
    value = result.get("token_spend")
    if isinstance(value, TokenSpend):
        return value
    if isinstance(value, dict):
        return TokenSpend.model_validate(value)
    return TokenSpend()


def _as_state(result: dict[str, Any]) -> RunState:
    """View a graph result as RunState, dropping LangGraph's internal keys."""
    cleaned = {key: value for key, value in result.items() if not key.startswith("__")}
    return cast(RunState, cleaned)


_service: RunService | None = None


def get_run_service() -> RunService:
    """Process-wide run service singleton."""
    global _service  # noqa: PLW0603 - deliberate process-wide singleton
    if _service is None:
        _service = RunService()
    return _service


def reset_run_service() -> None:
    """Test hook."""
    global _service  # noqa: PLW0603 - deliberate process-wide singleton
    _service = None