OrgTJWater/next-tjwater-server
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: simplify task and history services

Browse Source 
Remove retired query paths and repository compatibility branches, and share task test fixtures.
Cancel pending questions and resolve their notifications when recovering expired runs.
This commit is contained in:
jiang
2026-11-09 17:55:00 +08:00
parent 5197a93163
commit 6a4478bfe4
10 changed files with 464 additions and 816 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/services/operational_tasks.py
+65 -116
View File
@@ -3,7 +3,7 @@ from __future__ import annotations
import json
from datetime import datetime, timedelta, timezone
from typing import Iterable
from typing import TYPE_CHECKING, Iterable
from uuid import NAMESPACE_URL, UUID, uuid4, uuid5
from app.auth.permissions import (
@@ -77,6 +77,10 @@ from app.services.operational_scada_coverage import require_complete_scada_evide
from app.services.time_comparison_chart import validate_time_comparison_chart
if TYPE_CHECKING:
from app.infra.db.postgresql.operational_tasks import OperationalTaskRepository
LEASE_DURATION = timedelta(seconds=60)
@@ -89,7 +93,7 @@ def _task_json(model: OperationalTaskSummary) -> dict:
class OperationalTaskService:
def __init__(self, repository, context, permissions: Iterable[str], *, clock=None,
def __init__(self, repository: OperationalTaskRepository, context, permissions: Iterable[str], *, clock=None,
execution_scope: str = "production"):
self.repo = repository
self.context = context
@@ -114,42 +118,14 @@ class OperationalTaskService:
require(owner == command.agent_instance_id,
"AGENT_INSTANCE_MISMATCH", "Agent instance does not own this run", status=409)
@property
def _supports_runs(self) -> bool:
return hasattr(self.repo, "get_active_run") and hasattr(self.repo, "update_run")
async def _active_run(self, task: dict) -> dict | None:
if self._supports_runs:
return await self.repo.get_active_run(str(task["task_id"]))
if task.get("status") not in {"queued", "running"}:
return None
now = self._now()
return _json(OperationalTaskRun(
run_id=uuid5(NAMESPACE_URL, f"legacy-run:{task['task_id']}:{task['execution_generation']}"),
task_id=task["task_id"],
generation=task["execution_generation"],
execution_request_id=task["execution_request_id"],
status="running" if task["status"] == "running" else "queued",
revision=task["revision"],
agent_session_id=task.get("agent_session_id"),
lease_expires_at=task.get("agent_lease_expires_at"),
heartbeat_at=task.get("updated_at") if task["status"] == "running" else None,
started_at=task.get("created_at") if task["status"] == "running" else None,
dispatch_attempts=0,
next_dispatch_at=task.get("created_at") or now,
created_at=task.get("created_at") or now,
updated_at=task.get("updated_at") or now,
))
return await self.repo.get_active_run(str(task["task_id"]))
async def _append_event(self, **kwargs) -> int | None:
append_event = getattr(self.repo, "append_event", None)
return await append_event(**kwargs) if append_event is not None else None
async def _append_event(self, **kwargs) -> int:
return await self.repo.append_event(**kwargs)
async def _list_access_requests(self, task_id: UUID | str) -> list[dict]:
list_access_requests = getattr(self.repo, "list_access_requests", None)
if list_access_requests is None:
return []
return await list_access_requests(str(task_id))
return await self.repo.list_access_requests(str(task_id))
def _permission_view(self) -> OperationalTaskPermissions:
permissions = self.permissions
@@ -286,8 +262,7 @@ class OperationalTaskService:
actions = await self.repo.list_actions(str(task_id))
access_requests = await self._list_access_requests(task_id)
active_run = await self._active_run(task)
latest_run = (await self.repo.get_latest_run(str(task_id))
if hasattr(self.repo, "get_latest_run") else active_run)
latest_run = await self.repo.get_latest_run(str(task_id))
progress = await self.repo.get_latest_progress(str(task_id), task["execution_generation"])
work_order_snapshot = None
if task["task_type"] == "external_work_order" and task["source"].get("external_id"):
@@ -397,9 +372,8 @@ class OperationalTaskService:
replay = await self._request_replay(scope, command.request_id, body)
if replay:
return TaskAccessRequest.model_validate(replay["access_request"])
availability = getattr(self.repo, "access_requests_available", None)
require(
availability is None or await availability(),
await self.repo.access_requests_available(),
"DEPENDENCY_UNAVAILABLE",
"Task data access requests are not enabled for this project database",
status=503,
@@ -616,8 +590,7 @@ class OperationalTaskService:
return OperationalTaskRun.model_validate(await self._run(task_id, run_id))
async def _fail_run_messages(self, run_id, reason: str, now: datetime) -> None:
if hasattr(self.repo, "fail_run_messages"):
await self.repo.fail_run_messages(str(run_id), reason, now)
await self.repo.fail_run_messages(str(run_id), reason, now)
async def fail_run(
self, task_id: UUID, run_id: UUID, command: AgentRunFailureCommand
@@ -772,7 +745,7 @@ class OperationalTaskService:
# Lease recovery already holds a row lock. Try rather than wait
# for the shared run lock to avoid reversing the lock order of
# message/finish transactions; a busy run is checked next tick.
if hasattr(self.repo, "try_lock") and not await self.repo.try_lock(f"run:{run['run_id']}"):
if not await self.repo.try_lock(f"run:{run['run_id']}"):
continue
task_id = UUID(str(run["task_id"]))
expected = run["revision"]
@@ -784,6 +757,7 @@ class OperationalTaskService:
timed_out = OperationalTaskRun.model_validate(run)
await self.repo.update_run(_json(timed_out), expected)
await self._fail_run_messages(run["run_id"], "Run lease expired", now)
await task_questions.cancel_pending(self, task_id, now)
await self._run_event(timed_out, now)
task = await self._task(task_id)
timeout_count = await self.repo.consecutive_timed_out_runs(str(task_id))
@@ -1017,25 +991,9 @@ class OperationalTaskService:
lease_expires_at=now + LEASE_DURATION,
updated_at=now,
)
if self._supports_runs:
await self.repo.update_run(
_json(OperationalTaskRun.model_validate(active_run)), run_expected
)
else:
# Compatibility for pre-015 repository adapters only.
if command.report:
if task["task_type"] == "event_investigation":
task["status"] = self._investigation_status(actions.values(), task) if command.report.stage == "final" else "running"
else:
task["status"] = self._status_from_plan(current_plan.plan.nodes, actions.values()) if current_plan else "running"
if task["status"] in {"waiting_human", "completed", "failed", "cancelled"}:
task["lifecycle_status"] = task["status"]
task["agent_lease_expires_at"] = (
now + LEASE_DURATION if task["status"] == "running" else None
)
await self.repo.update_task(
_task_json(OperationalTaskSummary.model_validate(task))
)
await self.repo.update_run(
_json(OperationalTaskRun.model_validate(active_run)), run_expected
)
await self._task_event(task, now)
response = {"plan_revision": plan_revision, "report_revision": report_revision}
if progress_revision is not None:
@@ -1988,13 +1946,12 @@ class OperationalTaskService:
return DispatchDueSchedulesResult(task_ids=task_ids, dispatched_at=now)
async def _require_no_pending_calculations(self, task_id: UUID, run: dict) -> None:
if hasattr(self.repo, "pending_calculation_jobs"):
pending = await self.repo.pending_calculation_jobs(
str(task_id), str(run["run_id"]), str(run["execution_request_id"]),
)
require(not pending, "PENDING_CALCULATIONS",
"Wait for submitted calculations before publishing a final report or completing the run",
refs=pending, status=409, recovery="retry_same_request")
pending = await self.repo.pending_calculation_jobs(
str(task_id), str(run["run_id"]), str(run["execution_request_id"]),
)
require(not pending, "PENDING_CALCULATIONS",
"Wait for submitted calculations before publishing a final report or completing the run",
refs=pending, status=409, recovery="retry_same_request")
async def _finish_run(
self, task_id: UUID, run_id: UUID, command: AgentRunCommand,
@@ -2024,11 +1981,10 @@ class OperationalTaskService:
actions = []
if outcome == "succeeded":
await self._require_no_pending_calculations(task_id, run)
if hasattr(self.repo, "pending_run_messages"):
pending_messages = await self.repo.pending_run_messages(str(run_id))
require(not pending_messages, "PENDING_AGENT_MESSAGES",
"Run has messages not yet entered into Agent context",
status=409, recovery="retry_same_request", refs=pending_messages)
pending_messages = await self.repo.pending_run_messages(str(run_id))
require(not pending_messages, "PENDING_AGENT_MESSAGES",
"Run has messages not yet entered into Agent context",
status=409, recovery="retry_same_request", refs=pending_messages)
plan_document = await self.repo.get_latest_plan(str(task_id))
report_document = await self.repo.get_latest_report(str(task_id))
plan = (
@@ -2158,8 +2114,6 @@ class OperationalTaskService:
from app.infra.db.postgresql.task_alerts import TaskAlertRepository
from app.services.task_alerts import TaskAlertService
if not hasattr(self.repo, "list_investigation_alerts"):
return
for alert in await self.repo.list_investigation_alerts(task_id):
command = TaskAlertEvent.model_validate({
**{key: value for key, value in alert.items() if key not in {"status", "updated_at", "revision"}},
@@ -2364,42 +2318,41 @@ class OperationalTaskService:
async def _requeue_task(self, task: dict, now: datetime) -> None:
await task_questions.cancel_pending(self, UUID(str(task["task_id"])), now)
if self._supports_runs:
active_run = await self.repo.get_active_run(str(task["task_id"]))
while active_run is not None:
await self.repo.lock(f"run:{active_run['run_id']}")
# Completion may have won the run lock after our first read.
refreshed_run = await self.repo.get_active_run(str(task["task_id"]))
if refreshed_run is None or str(refreshed_run["run_id"]) == str(active_run["run_id"]):
active_run = refreshed_run
break
active_run = await self.repo.get_active_run(str(task["task_id"]))
while active_run is not None:
await self.repo.lock(f"run:{active_run['run_id']}")
# Completion may have won the run lock after our first read.
refreshed_run = await self.repo.get_active_run(str(task["task_id"]))
if refreshed_run is None or str(refreshed_run["run_id"]) == str(active_run["run_id"]):
active_run = refreshed_run
if active_run is not None and active_run["status"] == "queued":
# Multiple inputs can arrive before the worker claims the next
# generation. They are all visible when that queued run starts,
# so coalesce them instead of violating the one-active-run rule.
task.update(
lifecycle_status="open", status="queued",
revision=task["revision"] + 1,
execution_generation=active_run["generation"],
execution_request_id=str(active_run["execution_request_id"]),
agent_session_id=None, agent_lease_expires_at=None,
updated_at=now,
)
return
if active_run is not None:
expected = active_run["revision"]
active_run.update(
status="cancelled", revision=expected + 1,
lease_expires_at=None, finished_at=now,
failure_code="SUPERSEDED_BY_REQUEUE",
failure_detail="New task input superseded the active run",
updated_at=now,
)
cancelled = OperationalTaskRun.model_validate(active_run)
await self.repo.update_run(_json(cancelled), expected)
await self._fail_run_messages(active_run["run_id"], "Run superseded by requeue", now)
await self._run_event(cancelled, now)
break
active_run = refreshed_run
if active_run is not None and active_run["status"] == "queued":
# Multiple inputs can arrive before the worker claims the next
# generation. They are all visible when that queued run starts,
# so coalesce them instead of violating the one-active-run rule.
task.update(
lifecycle_status="open", status="queued",
revision=task["revision"] + 1,
execution_generation=active_run["generation"],
execution_request_id=str(active_run["execution_request_id"]),
agent_session_id=None, agent_lease_expires_at=None,
updated_at=now,
)
return
if active_run is not None:
expected = active_run["revision"]
active_run.update(
status="cancelled", revision=expected + 1,
lease_expires_at=None, finished_at=now,
failure_code="SUPERSEDED_BY_REQUEUE",
failure_detail="New task input superseded the active run",
updated_at=now,
)
cancelled = OperationalTaskRun.model_validate(active_run)
await self.repo.update_run(_json(cancelled), expected)
await self._fail_run_messages(active_run["run_id"], "Run superseded by requeue", now)
await self._run_event(cancelled, now)
generation = task["execution_generation"] + 1
task_id = UUID(task["task_id"])
task.update(
@@ -2409,9 +2362,8 @@ class OperationalTaskService:
agent_session_id=None, agent_lease_expires_at=None, updated_at=now,
)
run = self._new_run(task_id, generation, now)
if self._supports_runs:
await self.repo.insert_run(_json(run))
await self._run_event(run, now, queued=True)
await self.repo.insert_run(_json(run))
await self._run_event(run, now, queued=True)
async def _notify(
self, task_id: UUID, kind: str, title: str, description: str, now: datetime,
@@ -2447,10 +2399,7 @@ class OperationalTaskService:
now: datetime,
reason: str,
) -> None:
resolve = getattr(self.repo, "resolve_notifications", None)
if resolve is None:
return
notifications = await resolve(
notifications = await self.repo.resolve_notifications(
str(task_id),
kinds=kinds,
subject_type=subject_type,