Add Phase 2 durable workstation agent job routing
This commit is contained in:
A R R R Associates
@@ -0,0 +1,257 @@
|
||||
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())
|
||||
Reference in New Issue
Block a user