arrr-erp/app/modules/documents/agent_jobs.py

from __future__ import annotations

import json
from datetime import datetime, timedelta, timezone
from typing import Any
from uuid import uuid4

from sqlalchemy import or_, select
from sqlalchemy.exc import IntegrityError

from app.modules.documents.models import ERPAgentJob, ERPAgentJobEvent, ERPWorkstationAgent


LEASE_SECONDS = 45
DEFAULT_MAX_ATTEMPTS = 3
TERMINAL_STATUSES = {"succeeded", "failed", "cancelled"}


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


def _json_dump(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"), default=str)


def _json_load(value: str | None, default: Any) -> Any:
    if not value:
        return default
    try:
        parsed = json.loads(value)
        return parsed
    except Exception:
        return default


def _event(db, job: ERPAgentJob, event_type: str, *, detail: str | None = None, actor_user_id: int | None = None) -> None:
    db.add(
        ERPAgentJobEvent(
            job_id=job.id,
            tenant_id=job.tenant_id,
            branch_id=job.branch_id,
            workstation_agent_id=job.workstation_agent_id,
            event_type=str(event_type or "event")[:40],
            detail=(detail or None),
            actor_user_id=actor_user_id,
            occurred_at_utc=_utcnow(),
        )
    )


def enqueue_agent_job(
    db,
    *,
    workstation_agent_id: int,
    action: str,
    payload: dict[str, Any] | None = None,
    idempotency_key: str | None = None,
    priority: int = 5,
    max_attempts: int = DEFAULT_MAX_ATTEMPTS,
    created_by_user_id: int | None = None,
) -> ERPAgentJob:
    workstation = db.get(ERPWorkstationAgent, int(workstation_agent_id))
    if not workstation or not workstation.is_active:
        raise ValueError("The selected ERP workstation agent is unavailable or inactive.")
    action_value = str(action or "").strip()
    if not action_value:
        raise ValueError("Agent job action is required.")
    idem = str(idempotency_key or "").strip() or None
    if idem:
        existing = db.execute(
            select(ERPAgentJob).where(
                ERPAgentJob.workstation_agent_id == workstation.id,
                ERPAgentJob.idempotency_key == idem,
            )
        ).scalar_one_or_none()
        if existing:
            return existing
    job = ERPAgentJob(
        job_uuid=uuid4().hex,
        tenant_id=workstation.tenant_id,
        branch_id=workstation.branch_id,
        storage_node_id=workstation.storage_node_id,
        workstation_agent_id=workstation.id,
        action=action_value[:120],
        payload_json=_json_dump(payload or {}),
        idempotency_key=idem[:200] if idem else None,
        status="queued",
        priority=max(0, min(int(priority), 100)),
        max_attempts=max(1, min(int(max_attempts), 20)),
        attempts=0,
        created_by_user_id=created_by_user_id,
        created_at_utc=_utcnow(),
        updated_at_utc=_utcnow(),
    )
    db.add(job)
    try:
        db.flush()
    except IntegrityError:
        db.rollback()
        if idem:
            existing = db.execute(
                select(ERPAgentJob).where(
                    ERPAgentJob.workstation_agent_id == workstation.id,
                    ERPAgentJob.idempotency_key == idem,
                )
            ).scalar_one_or_none()
            if existing:
                return existing
        raise
    _event(db, job, "queued", actor_user_id=created_by_user_id)
    return job


def claim_jobs_for_workstation(db, *, storage_node_id: int, agent_instance_id: str, limit: int = 10) -> list[dict[str, Any]]:
    instance_id = str(agent_instance_id or "").strip()
    if not instance_id:
        return []
    workstation = db.execute(
        select(ERPWorkstationAgent).where(
            ERPWorkstationAgent.storage_node_id == int(storage_node_id),
            ERPWorkstationAgent.agent_instance_id == instance_id,
            ERPWorkstationAgent.is_active.is_(True),
        )
    ).scalar_one_or_none()
    if not workstation:
        return []

    now = _utcnow()
    # Expired claims become retryable. Jobs that have exhausted their attempts fail.
    expired = list(
        db.execute(
            select(ERPAgentJob).where(
                ERPAgentJob.workstation_agent_id == workstation.id,
                ERPAgentJob.status == "claimed",
                ERPAgentJob.lease_expires_at_utc.is_not(None),
                ERPAgentJob.lease_expires_at_utc < now,
            )
        ).scalars().all()
    )
    for job in expired:
        if int(job.attempts or 0) >= int(job.max_attempts or DEFAULT_MAX_ATTEMPTS):
            job.status = "failed"
            job.failed_at_utc = now
            job.last_error = job.last_error or "Agent job lease expired after maximum retry attempts."
            _event(db, job, "failed", detail=job.last_error)
        else:
            job.status = "queued"
            job.claimed_at_utc = None
            job.lease_expires_at_utc = None
            _event(db, job, "lease_expired", detail="Job returned to queue for retry.")

    stmt = (
        select(ERPAgentJob)
        .where(
            ERPAgentJob.workstation_agent_id == workstation.id,
            ERPAgentJob.status == "queued",
            ERPAgentJob.attempts < ERPAgentJob.max_attempts,
        )
        .order_by(ERPAgentJob.priority.desc(), ERPAgentJob.created_at_utc.asc(), ERPAgentJob.id.asc())
        .limit(max(1, min(int(limit), 25)))
    )
    rows = list(db.execute(stmt).scalars().all())
    payloads: list[dict[str, Any]] = []
    for job in rows:
        job.status = "claimed"
        job.attempts = int(job.attempts or 0) + 1
        job.claimed_at_utc = now
        job.lease_expires_at_utc = now + timedelta(seconds=LEASE_SECONDS)
        job.updated_at_utc = now
        _event(db, job, "claimed", detail=f"attempt={job.attempts}")
        payloads.append(
            {
                "job_uuid": job.job_uuid,
                "action": job.action,
                "payload": _json_load(job.payload_json, {}),
                "idempotency_key": job.idempotency_key or job.job_uuid,
                "attempt": int(job.attempts or 0),
                "max_attempts": int(job.max_attempts or DEFAULT_MAX_ATTEMPTS),
                "lease_seconds": LEASE_SECONDS,
            }
        )
    return payloads


def complete_agent_job(
    db,
    *,
    storage_node_id: int,
    agent_instance_id: str,
    message: dict[str, Any],
) -> ERPAgentJob | None:
    instance_id = str(agent_instance_id or "").strip()
    job_uuid = str(message.get("job_uuid") or "").strip()
    if not instance_id or not job_uuid:
        return None
    workstation = db.execute(
        select(ERPWorkstationAgent).where(
            ERPWorkstationAgent.storage_node_id == int(storage_node_id),
            ERPWorkstationAgent.agent_instance_id == instance_id,
            ERPWorkstationAgent.is_active.is_(True),
        )
    ).scalar_one_or_none()
    if not workstation:
        return None
    job = db.execute(
        select(ERPAgentJob).where(
            ERPAgentJob.job_uuid == job_uuid,
            ERPAgentJob.workstation_agent_id == workstation.id,
        )
    ).scalar_one_or_none()
    if not job:
        return None
    if job.status in TERMINAL_STATUSES:
        return job

    now = _utcnow()
    ok = bool(message.get("ok"))
    job.result_json = _json_dump(message.get("result") or {}) if ok else None
    job.last_error = None if ok else str(message.get("error") or "Agent job failed.")[:4000]
    job.agent_completed_at_utc = now
    job.updated_at_utc = now
    job.lease_expires_at_utc = None
    if ok:
        job.status = "succeeded"
        job.completed_at_utc = now
        _event(db, job, "succeeded")
    elif int(job.attempts or 0) < int(job.max_attempts or DEFAULT_MAX_ATTEMPTS) and bool(message.get("retryable", False)):
        job.status = "queued"
        job.claimed_at_utc = None
        _event(db, job, "retry_queued", detail=job.last_error)
    else:
        job.status = "failed"
        job.failed_at_utc = now
        _event(db, job, "failed", detail=job.last_error)
    return job


def cancel_agent_job(db, job: ERPAgentJob, *, actor_user_id: int | None = None, reason: str | None = None) -> ERPAgentJob:
    if job.status in TERMINAL_STATUSES:
        return job
    job.status = "cancelled"
    job.cancelled_at_utc = _utcnow()
    job.updated_at_utc = job.cancelled_at_utc
    job.lease_expires_at_utc = None
    _event(db, job, "cancelled", detail=reason, actor_user_id=actor_user_id)
    return job


def list_recent_agent_jobs(db, *, tenant_id: int | None = None, branch_id: int | None = None, limit: int = 25) -> list[ERPAgentJob]:
    stmt = select(ERPAgentJob).order_by(ERPAgentJob.created_at_utc.desc(), ERPAgentJob.id.desc())
    if tenant_id is not None:
        stmt = stmt.where(ERPAgentJob.tenant_id == int(tenant_id))
    if branch_id is not None:
        stmt = stmt.where(ERPAgentJob.branch_id == int(branch_id))
    stmt = stmt.limit(max(1, min(int(limit), 200)))
    return list(db.execute(stmt).scalars().all())