jiang
6a4478bfe4
Remove retired query paths and repository compatibility branches, and share task test fixtures. Cancel pending questions and resolve their notifications when recovering expired runs.
253 lines
9.4 KiB
Python
253 lines
9.4 KiB
Python
"""In-memory task/run fixtures exercising the production repository contract."""
|
|
from contextlib import asynccontextmanager
|
|
from copy import deepcopy
|
|
from datetime import datetime, timedelta
|
|
from uuid import uuid4
|
|
from app.auth.project_dependencies import ProjectContext
|
|
from app.domain.schemas.operational_task import OperationalTaskSummary, OperationalTaskRun
|
|
|
|
def _context(user_id=None, *, system_role="user"):
|
|
return ProjectContext(
|
|
project_id=uuid4(), project_code="test", user_id=user_id or uuid4(),
|
|
project_role="member", system_role=system_role,
|
|
)
|
|
|
|
def _task_document(context, now, *, lease=None):
|
|
task_id = uuid4()
|
|
return OperationalTaskSummary(
|
|
execution_scope="production",
|
|
project_id=context.project_id,
|
|
task_id=task_id,
|
|
task_type="event_investigation",
|
|
workflow="water_quality_complaint",
|
|
status="running",
|
|
priority="normal",
|
|
revision=2,
|
|
execution_generation=1,
|
|
execution_request_id=uuid4(),
|
|
title="调查",
|
|
summary="验证状态机",
|
|
source={"kind": "operator", "event_kind": "water_quality_complaint"},
|
|
trigger={"kind": "manual", "occurred_at": now, "detail": "test"},
|
|
scheduled_for=None,
|
|
agent_session_id="session-1",
|
|
agent_lease_expires_at=lease,
|
|
current_plan_revision=None,
|
|
current_report_revision=None,
|
|
created_at=now,
|
|
updated_at=now,
|
|
).model_dump(mode="json")
|
|
|
|
def _active_run(task, now):
|
|
return OperationalTaskRun(
|
|
run_id=uuid4(), task_id=task["task_id"], generation=task["execution_generation"],
|
|
execution_request_id=task["execution_request_id"], status="running", revision=2,
|
|
agent_session_id="session-1", lease_expires_at=now + timedelta(minutes=1),
|
|
heartbeat_at=now, started_at=now, dispatch_attempts=0,
|
|
next_dispatch_at=now, created_at=now, updated_at=now,
|
|
).model_dump(mode="json")
|
|
|
|
class TaskRepository:
|
|
def __init__(self, task, run=None, *, action=None):
|
|
self.task = task
|
|
self.action = action
|
|
self.requests = {}
|
|
self.evidence = {}
|
|
self.feedback = {}
|
|
self.plans = []
|
|
self.reports = []
|
|
self.progress = {}
|
|
self.notifications = []
|
|
self.locks = []
|
|
self.access_requests = {}
|
|
self.events = []
|
|
self.questions = {}
|
|
now = datetime.fromisoformat(task["created_at"])
|
|
seeded = _active_run(task, now) if run is None else run
|
|
if run is None:
|
|
seeded.update(
|
|
status=task["status"] if task["status"] in {"queued", "running"} else "succeeded",
|
|
agent_session_id=task.get("agent_session_id"),
|
|
lease_expires_at=task.get("agent_lease_expires_at"),
|
|
)
|
|
self.run = seeded
|
|
self.runs = {seeded["run_id"]: seeded}
|
|
|
|
@asynccontextmanager
|
|
async def transaction(self):
|
|
yield self
|
|
|
|
async def lock(self, key):
|
|
self.locks.append(key)
|
|
|
|
async def get_request(self, scope, request_id):
|
|
return self.requests.get((scope, request_id))
|
|
|
|
async def save_request(self, scope, request_id, request_digest, response):
|
|
self.requests[(scope, request_id)] = {
|
|
"request_digest": request_digest,
|
|
"response": response,
|
|
}
|
|
|
|
async def get_task(self, task_id):
|
|
return self.task if task_id == self.task["task_id"] else None
|
|
|
|
async def update_task(self, task):
|
|
self.task = task
|
|
|
|
async def get_latest_plan(self, _task_id):
|
|
return self.plans[-1] if self.plans else None
|
|
|
|
async def get_latest_report(self, _task_id):
|
|
return self.reports[-1] if self.reports else None
|
|
|
|
async def get_latest_progress(self, task_id, generation):
|
|
return self.progress.get((task_id, generation))
|
|
|
|
async def save_progress(self, task_id, snapshot):
|
|
self.progress[(task_id, snapshot["execution_generation"])] = snapshot
|
|
|
|
async def get_report(self, _task_id, revision):
|
|
return next((item for item in self.reports if item["revision"] == revision), None)
|
|
|
|
async def list_reports(self, _task_id, limit, offset):
|
|
return list(reversed(self.reports))[offset:offset + limit]
|
|
|
|
async def insert_plan(self, _task_id, snapshot):
|
|
self.plans.append(snapshot)
|
|
|
|
async def insert_report(self, _task_id, snapshot):
|
|
self.reports.append(snapshot)
|
|
|
|
async def list_evidence(self, _task_id):
|
|
return list(self.evidence.values())
|
|
|
|
async def get_evidence(self, _task_id, evidence_id):
|
|
return self.evidence.get(evidence_id)
|
|
|
|
async def insert_evidence(self, _task_id, evidence):
|
|
self.evidence[evidence["evidence_id"]] = evidence
|
|
|
|
async def list_access_requests(self, _task_id):
|
|
return list(self.access_requests.values())
|
|
|
|
async def get_access_request(self, _task_id, access_request_id):
|
|
return self.access_requests.get(access_request_id)
|
|
|
|
async def insert_access_request(self, _task_id, document):
|
|
self.access_requests[document["access_request_id"]] = document
|
|
|
|
async def update_access_request(self, _task_id, document, expected_version):
|
|
assert document["version"] == expected_version + 1
|
|
self.access_requests[document["access_request_id"]] = document
|
|
|
|
async def list_actions(self, _task_id):
|
|
return [self.action] if self.action else []
|
|
|
|
async def get_action(self, _task_id, action_id):
|
|
if self.action and self.action["action_id"] == action_id:
|
|
return self.action
|
|
return None
|
|
|
|
async def update_action(self, _task_id, action):
|
|
self.action = action
|
|
|
|
async def insert_action(self, _task_id, action):
|
|
self.action = action
|
|
|
|
async def insert_notification(self, notification):
|
|
self.notifications.append(notification)
|
|
|
|
async def resolve_notifications(
|
|
self, task_id, *, kinds, subject_type, subject_id, resolved_at, resolution_reason
|
|
):
|
|
resolved = []
|
|
for notification in self.notifications:
|
|
if (
|
|
notification["task_id"] == task_id
|
|
and notification["kind"] in kinds
|
|
and notification.get("lifecycle_status", "active") == "active"
|
|
and notification.get("subject_type") == subject_type
|
|
and notification.get("subject_id") == subject_id
|
|
):
|
|
notification.update(
|
|
lifecycle_status="resolved",
|
|
resolved_at=resolved_at.isoformat(),
|
|
resolution_reason=resolution_reason,
|
|
)
|
|
resolved.append(notification.copy())
|
|
return resolved
|
|
|
|
async def get_feedback(self, task_id, source_event_id):
|
|
return self.feedback.get((task_id, source_event_id))
|
|
|
|
async def upsert_feedback(self, task_id, feedback, accepted_at):
|
|
self.feedback[(task_id, feedback["source_event_id"])] = {
|
|
"source_revision": feedback["source_revision"],
|
|
"document": feedback,
|
|
"accepted_at": accepted_at,
|
|
}
|
|
|
|
async def get_active_run(self, task_id):
|
|
return next((run for run in sorted(self.runs.values(), key=lambda item: item["generation"], reverse=True)
|
|
if run["task_id"] == task_id and run["status"] in {"queued", "claimed", "running"}), None)
|
|
|
|
async def get_latest_run(self, task_id):
|
|
return next((run for run in sorted(self.runs.values(), key=lambda item: item["generation"], reverse=True) if run["task_id"] == task_id), None)
|
|
|
|
async def get_run(self, run_id):
|
|
return self.runs.get(run_id)
|
|
|
|
async def require_execution_storage(self):
|
|
return None
|
|
|
|
async def update_run(self, run, expected_revision):
|
|
assert run["revision"] == expected_revision + 1
|
|
self.runs[run["run_id"]] = run
|
|
self.run = max(self.runs.values(), key=lambda item: item["generation"])
|
|
|
|
async def insert_run(self, run):
|
|
self.runs[run["run_id"]] = run
|
|
self.run = max(self.runs.values(), key=lambda item: item["generation"])
|
|
|
|
async def append_event(self, **event):
|
|
self.events.append(event)
|
|
return len(self.events)
|
|
|
|
async def current_event_cursor(self):
|
|
return len(self.events)
|
|
|
|
async def list_presentation_events(self, task_id, after, through, limit):
|
|
return [{"event_seq": index, **event} for index, event in enumerate(self.events, 1)
|
|
if after < index <= through and event["topic"] == "operational.presentation.requested"
|
|
and event["aggregate_id"] == task_id][:limit]
|
|
async def access_requests_available(self):
|
|
return True
|
|
|
|
async def fail_run_messages(self, run_id, reason, now):
|
|
return None
|
|
|
|
async def try_lock(self, key):
|
|
await self.lock(key)
|
|
return True
|
|
|
|
async def pending_calculation_jobs(self, task_id, run_id, execution_request_id):
|
|
return []
|
|
|
|
async def pending_run_messages(self, run_id):
|
|
return []
|
|
|
|
async def list_investigation_alerts(self, task_id):
|
|
return []
|
|
|
|
async def list_questions(self, task_id):
|
|
return [deepcopy(q) for q in self.questions.values() if q["task_id"] == task_id]
|
|
|
|
async def get_question(self, task_id, question_id):
|
|
q = self.questions.get(question_id)
|
|
return deepcopy(q) if q and q["task_id"] == task_id else None
|
|
|
|
async def save_question(self, task_id, document):
|
|
assert document["task_id"] == task_id
|
|
self.questions[document["question_request_id"]] = deepcopy(document)
|