PTSRaiAgent/app/scheduler/jobs.py

"""Questionnaire jobs on the demo clock; clinical release disabled by default.

One scheduler per database. PostgreSQL's advisory transaction lock serializes
workers; an asyncio lock does the same within a SQLite demonstration process.
Deliveries are at least once: the channel protocol offers no idempotency key.
"""

import asyncio
import logging
from contextlib import suppress
from datetime import datetime
from zoneinfo import ZoneInfo

from sqlalchemy import select, text

from app.core import clock
from app.core.intake_runtime import due_scales, intake_enabled, load_runtime
from app.db import repo
from app.db.clinical import PatientProfile
from app.db.intake import QuestionnaireTask
from app.db.session import get_clinical_session
from app.dialog.contracts import Event, Reply, get_engine
from app.integrations.base import integrations
from app.services import questionnaire as service

log = logging.getLogger(__name__)
_worker = None
_lock = asyncio.Lock()


async def plan():
    now = clock.now()
    p = load_runtime()
    local = now.astimezone(ZoneInfo(p.timezone))
    async with get_clinical_session() as session:
        profiles = (await session.scalars(select(PatientProfile))).all()
        for profile in profiles:
            state = await repo.get_dialog_state(session, profile.token)
            if not state or state.data.get("consent") is not True or state.data.get("adult") is not True:
                continue
            rows = await service.records(session, profile.token)
            initial = next((r for r in reversed(rows) if r.kind == "intake" and r.status == "complete"), None)
            if initial is None:
                continue
            last_activity = max([r.day for r in rows if r.kind == "diary" and r.status == "complete"] + [initial.day])
            high_risk = profile.enhanced_monitoring or bool(state.data.get("crisis"))
            if high_risk and (local.date() - last_activity).days >= p.high_risk_silence_days:
                await service.alert(session, profile.token, f"silence:{last_activity}", "high_risk_silence", now)
            if (state.data.get("crisis") or {}).get("open") or state.data.get("intake_review_required"):
                continue
            if local.hour >= p.daily_hour and not any(
                r.kind == "diary" and r.day == local.date() and r.status == "complete" for r in rows
            ):
                await service.queue_task(session, profile.token, "diary", f"diary:{local.date()}", now)
            last = {}
            for row in rows:
                if row.status != "complete":
                    continue
                for sid in ("phq9", "pcl5", "rus_audit_s"):
                    if row.details.get("totals", {}).get(sid, {}).get("complete"):
                        last[sid] = row.completed_at.astimezone(ZoneInfo(p.timezone)).date()
            route = initial.details["decision"]["primary_route"]
            prescriptions = await integrations.prescriptions.get_prescriptions(profile.token)
            changes = [
                v.start_date
                for v in prescriptions
                if v.start_date and v.start_date <= local.date() and (v.end_date is None or v.end_date >= local.date())
            ]
            if profile.therapy_start and profile.therapy_start <= local.date():
                changes.append(profile.therapy_start)
            closed = (state.data.get("crisis") or {}).get("closed_at")
            crisis_date = datetime.fromisoformat(closed).astimezone(ZoneInfo(p.timezone)).date() if closed else None
            scales = due_scales(
                today=local.date(),
                enrolled=initial.day,
                last=last,
                route="self_help" if route == "self_help" else "specialist",
                medication_changed=max(changes) if changes else None,
                after_crisis=crisis_date,
                policy=p,
            )
            if scales or local.hour >= p.daily_hour:
                await service.queue_task(session, profile.token, "cssrs_since_contact", f"cssrs:{local.date()}", now)
            for sid in scales:
                if sid == "pcl5" and initial.answers.get("trauma") != 1:
                    continue
                await service.queue_task(session, profile.token, sid, f"{sid}:{local.date()}", now)


async def deliver():
    async with get_clinical_session() as session:
        task_ids = list(
            (
                await session.scalars(
                    select(QuestionnaireTask.id)
                    .where(
                        QuestionnaireTask.sent_at.is_(None),
                        QuestionnaireTask.status != "cancelled",
                        QuestionnaireTask.due_at <= clock.now(),
                    )
                    .order_by((QuestionnaireTask.kind != "cssrs_since_contact"), QuestionnaireTask.id)
                )
            ).all()
        )
    for task_id in task_ids:
        async with get_clinical_session() as session:
            task = await session.get(QuestionnaireTask, task_id)
            state = await repo.get_dialog_state(session, task.token)
            if not state or state.data.get("consent") is not True or state.data.get("adult") is not True:
                task.status, task.delivery = "cancelled", []
                continue
            if (state.data.get("crisis") or {}).get("open"):
                continue
            if task.status == "done" and not task.delivery:
                continue
            token = task.token
            payload = task.delivery
            if not payload and state.state not in ("idle", "intake.collect"):
                continue
        if not payload:
            await get_engine().handle(
                token, Event(type="system", system="daily_checkin", params={"questionnaire_task": task_id})
            )
        # Serialize delivery with withdrawal of consent and user answers.
        async with get_clinical_session() as session:
            await service.lock_patient(session, token)
            state = await repo.get_dialog_state(session, token)
            task = await session.get(QuestionnaireTask, task_id)
            if not state or state.data.get("consent") is not True or not task.delivery:
                continue
            try:
                await asyncio.wait_for(
                    integrations.outbox.push(token, [Reply.model_validate(r) for r in task.delivery]), timeout=10
                )
            except Exception as exc:
                log.warning("questionnaire outbox error=%s", type(exc).__name__)
            else:
                task.sent_at = clock.now()


async def run_until_now():
    async with _lock, get_clinical_session() as guard:
        if guard.bind.dialect.name == "postgresql" and not await guard.scalar(
            text("SELECT pg_try_advisory_xact_lock(520053)")
        ):
            return
        await service.dispatch_notices()
        if not intake_enabled():
            return
        await plan()
        await deliver()
        await service.dispatch_notices()


async def _run():
    while True:
        try:
            await run_until_now()
        except Exception as exc:
            log.error("questionnaire scheduler error=%s", type(exc).__name__)
        await asyncio.sleep(30)


async def start():
    global _worker
    if _worker is None and intake_enabled():
        _worker = asyncio.create_task(_run())


async def stop():
    global _worker
    if _worker is not None:
        _worker.cancel()
        with suppress(asyncio.CancelledError):
            await _worker
        _worker = None