Initial commit
This commit is contained in:
LZH-YS1998
@@ -0,0 +1,552 @@
|
||||
"""Pluggable work-item crash-recovery manager.
|
||||
|
||||
Detects interrupted company-mode work items after server restart and provides
|
||||
deterministic Resume / Cancel operations — no LLM needed.
|
||||
|
||||
Hooks into existing engine capabilities without modifying core code:
|
||||
- engine.store → task queries + persistence
|
||||
- engine._load_company_runtime_snapshot() → reconstruct work-item plan + tasks
|
||||
- engine.company_executor.execute() → resume execution
|
||||
- ws_handler.broadcast() → push status to all clients
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any, Awaitable, Callable
|
||||
|
||||
from opc.layer2_organization.work_item_identity import work_item_projection_id_from_metadata
|
||||
from opc.layer2_organization.work_item_transition import apply_task_status_transition
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Data types
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class RecoverableWorkItem:
|
||||
work_item_projection_id: str
|
||||
title: str
|
||||
task_id: str
|
||||
status: str # "done" | "failed" | "pending" | "blocked" | "cancelled"
|
||||
interrupted: bool # has interrupted_recovery metadata
|
||||
previous_status: str # what it was before reconciliation
|
||||
|
||||
|
||||
@dataclass
|
||||
class InterruptedWorkItemRun:
|
||||
parent_session_id: str
|
||||
parent_task_id: str
|
||||
project_id: str
|
||||
title: str
|
||||
profile: str
|
||||
interrupted_at: str
|
||||
work_items: list[RecoverableWorkItem] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RecoveryStatus:
|
||||
interrupted: list[InterruptedWorkItemRun] = field(default_factory=list)
|
||||
active_recoveries: list[str] = field(default_factory=list)
|
||||
scanned_at: float = 0.0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Manager
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class RuntimeRecoveryManager:
|
||||
"""Scans for interrupted work-item runs and provides deterministic recovery."""
|
||||
|
||||
_CACHE_TTL = 10.0 # seconds
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
engine: Any,
|
||||
broadcast_fn: Callable[[dict[str, Any]], Awaitable[None]],
|
||||
) -> None:
|
||||
self._engine = engine
|
||||
self._broadcast = broadcast_fn
|
||||
self._lock = asyncio.Lock()
|
||||
self._active_recoveries: dict[str, asyncio.Task[Any]] = {}
|
||||
self._cached: RecoveryStatus | None = None
|
||||
self._cache_until: float = 0.0
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def get_recovery_status(self) -> RecoveryStatus:
|
||||
"""Return cached scan results, re-scanning if stale."""
|
||||
now = time.time()
|
||||
if self._cached is not None and now < self._cache_until:
|
||||
# Keep active_recoveries up to date even from cache
|
||||
self._cached.active_recoveries = list(self._active_recoveries.keys())
|
||||
return self._cached
|
||||
status = await self.scan()
|
||||
self._cached = status
|
||||
self._cache_until = now + self._CACHE_TTL
|
||||
return status
|
||||
|
||||
async def scan(self) -> RecoveryStatus:
|
||||
"""Scan the database for recoverable interrupted work-item runs."""
|
||||
store = self._engine.store
|
||||
if not store:
|
||||
return RecoveryStatus()
|
||||
|
||||
project_id = self._engine.project_id or "default"
|
||||
try:
|
||||
all_tasks = await store.get_tasks(project_id=project_id)
|
||||
except Exception as exc:
|
||||
logger.warning(f"Recovery scan failed: {exc}")
|
||||
return RecoveryStatus()
|
||||
|
||||
# Group projected work-item tasks by parent_session_id.
|
||||
groups: dict[str, list[Any]] = {}
|
||||
all_tasks_by_session: dict[str, Any] = {}
|
||||
for task in all_tasks:
|
||||
session_id = str(getattr(task, "session_id", "") or "").strip()
|
||||
if session_id:
|
||||
all_tasks_by_session[session_id] = task
|
||||
|
||||
parent_sid = str(getattr(task, "parent_session_id", "") or "").strip()
|
||||
projection_id = work_item_projection_id_from_metadata(getattr(task, "metadata", {}) or {})
|
||||
if parent_sid and projection_id:
|
||||
groups.setdefault(parent_sid, []).append(task)
|
||||
|
||||
interrupted: list[InterruptedWorkItemRun] = []
|
||||
|
||||
for parent_sid, tasks in groups.items():
|
||||
# Check if any task has interrupted_recovery metadata
|
||||
has_interrupted = any(
|
||||
_is_interrupted(t) for t in tasks
|
||||
)
|
||||
if not has_interrupted:
|
||||
continue
|
||||
|
||||
# Skip if all tasks are terminal (DONE or CANCELLED)
|
||||
from opc.core.models import TaskStatus
|
||||
non_terminal = [
|
||||
t for t in tasks
|
||||
if t.status not in (TaskStatus.DONE, TaskStatus.CANCELLED)
|
||||
]
|
||||
if not non_terminal:
|
||||
continue
|
||||
|
||||
# Find the parent (primary) task
|
||||
parent_task = all_tasks_by_session.get(parent_sid)
|
||||
parent_task_id = parent_task.id if parent_task else parent_sid
|
||||
title = parent_task.title if parent_task else "Unknown work-item run"
|
||||
|
||||
# Build work-item list.
|
||||
work_items: list[RecoverableWorkItem] = []
|
||||
earliest_interrupt = ""
|
||||
for t in sorted(tasks, key=lambda x: (x.created_at, x.id)):
|
||||
meta = dict(getattr(t, "metadata", {}) or {})
|
||||
recovery_meta = meta.get("interrupted_recovery", {})
|
||||
is_int = _is_interrupted(t)
|
||||
if is_int and recovery_meta.get("detected_at", ""):
|
||||
detected = recovery_meta["detected_at"]
|
||||
if not earliest_interrupt or detected < earliest_interrupt:
|
||||
earliest_interrupt = detected
|
||||
|
||||
work_items.append(RecoverableWorkItem(
|
||||
work_item_projection_id=work_item_projection_id_from_metadata(meta, fallback=t.id),
|
||||
title=t.title,
|
||||
task_id=t.id,
|
||||
status=t.status.value if hasattr(t.status, "value") else str(t.status),
|
||||
interrupted=is_int,
|
||||
previous_status=recovery_meta.get("previous_status", ""),
|
||||
))
|
||||
|
||||
profile = ""
|
||||
for t in tasks:
|
||||
p = (getattr(t, "metadata", {}) or {}).get("company_profile", "")
|
||||
if p:
|
||||
profile = p
|
||||
break
|
||||
|
||||
interrupted.append(InterruptedWorkItemRun(
|
||||
parent_session_id=parent_sid,
|
||||
parent_task_id=parent_task_id,
|
||||
project_id=project_id,
|
||||
title=title,
|
||||
profile=profile,
|
||||
interrupted_at=earliest_interrupt or datetime.now().isoformat(),
|
||||
work_items=work_items,
|
||||
))
|
||||
|
||||
return RecoveryStatus(
|
||||
interrupted=interrupted,
|
||||
active_recoveries=list(self._active_recoveries.keys()),
|
||||
scanned_at=time.time(),
|
||||
)
|
||||
|
||||
async def resume(self, parent_task_id: str) -> dict[str, Any]:
|
||||
"""Deterministically resume an interrupted work-item run."""
|
||||
async with self._lock:
|
||||
if parent_task_id in self._active_recoveries:
|
||||
return {"ok": False, "error": "already_in_progress"}
|
||||
|
||||
# Find the interrupted work-item run.
|
||||
status = await self.scan()
|
||||
wf = next((w for w in status.interrupted if w.parent_task_id == parent_task_id), None)
|
||||
if not wf:
|
||||
return {"ok": False, "error": "not_found"}
|
||||
|
||||
# Load snapshot
|
||||
snapshot = await self._engine._load_company_runtime_snapshot(wf.parent_session_id)
|
||||
if not snapshot:
|
||||
return {"ok": False, "error": "snapshot_unavailable"}
|
||||
|
||||
plan, tasks = snapshot
|
||||
|
||||
# Clean orphaned checkpoints
|
||||
await self._clean_orphaned_checkpoints(wf, tasks)
|
||||
|
||||
# Reset interrupted/failed/blocked tasks → PENDING
|
||||
from opc.core.models import TaskStatus
|
||||
resumed_ids: list[str] = []
|
||||
failed_ids: list[str] = []
|
||||
for task in tasks:
|
||||
if task.status == TaskStatus.DONE:
|
||||
continue
|
||||
if task.status in (TaskStatus.FAILED, TaskStatus.BLOCKED):
|
||||
task.result = None
|
||||
task.execution_lock = False
|
||||
task.execution_locked_at = None
|
||||
meta = dict(task.metadata)
|
||||
meta.pop("interrupted_recovery", None)
|
||||
progress = list(meta.get("progress_log", []))
|
||||
progress.append(f"[Recovery] Resumed at {datetime.now().isoformat()}")
|
||||
meta["progress_log"] = progress[-20:]
|
||||
task.metadata = meta
|
||||
projection_id = work_item_projection_id_from_metadata(meta, fallback=task.id)
|
||||
try:
|
||||
await apply_task_status_transition(
|
||||
self._engine.store,
|
||||
task,
|
||||
target_status_or_phase=TaskStatus.PENDING,
|
||||
reason="office_recovery_resume",
|
||||
release_claim=True,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Recovery resume skipped %s: %s", task.id, exc)
|
||||
failed_ids.append(projection_id)
|
||||
continue
|
||||
if task.status != TaskStatus.PENDING:
|
||||
logger.warning("Recovery resume preserved non-runnable phase for %s", task.id)
|
||||
failed_ids.append(projection_id)
|
||||
continue
|
||||
await self._engine.store.save_task(task)
|
||||
resumed_ids.append(projection_id)
|
||||
|
||||
if not resumed_ids:
|
||||
return {
|
||||
"ok": False,
|
||||
"error": "no_work_items_to_resume",
|
||||
"failed_work_item_projection_ids": failed_ids,
|
||||
}
|
||||
|
||||
# Invalidate cache
|
||||
self._cache_until = 0.0
|
||||
|
||||
# Launch execution in background
|
||||
await self._set_run_recovery_state(
|
||||
wf.parent_session_id,
|
||||
status="resuming",
|
||||
lifecycle_status="active",
|
||||
)
|
||||
bg_task = asyncio.create_task(
|
||||
self._execute_recovery(parent_task_id, wf.parent_session_id, plan, tasks)
|
||||
)
|
||||
self._active_recoveries[parent_task_id] = bg_task
|
||||
|
||||
return {
|
||||
"ok": True,
|
||||
"resumed_work_item_projection_ids": resumed_ids,
|
||||
"failed_work_item_projection_ids": failed_ids,
|
||||
}
|
||||
|
||||
async def cancel(self, parent_task_id: str) -> dict[str, Any]:
|
||||
"""Cancel an interrupted work-item run and clean up."""
|
||||
async with self._lock:
|
||||
# Cancel active recovery if running
|
||||
bg = self._active_recoveries.pop(parent_task_id, None)
|
||||
if bg and not bg.done():
|
||||
bg.cancel()
|
||||
|
||||
# Find the interrupted work-item run.
|
||||
status = await self.scan()
|
||||
wf = next((w for w in status.interrupted if w.parent_task_id == parent_task_id), None)
|
||||
if not wf:
|
||||
return {"ok": False, "error": "not_found"}
|
||||
|
||||
# Load tasks and cancel non-terminal ones
|
||||
snapshot = await self._engine._load_company_runtime_snapshot(wf.parent_session_id)
|
||||
if not snapshot:
|
||||
return {"ok": False, "error": "snapshot_unavailable"}
|
||||
|
||||
_, tasks = snapshot
|
||||
from opc.core.models import TaskStatus
|
||||
cancelled_count = 0
|
||||
failed_ids: list[str] = []
|
||||
for task in tasks:
|
||||
if task.status not in (TaskStatus.DONE, TaskStatus.CANCELLED):
|
||||
projection_id = work_item_projection_id_from_metadata(
|
||||
getattr(task, "metadata", {}) or {},
|
||||
fallback=task.id,
|
||||
)
|
||||
try:
|
||||
await apply_task_status_transition(
|
||||
self._engine.store,
|
||||
task,
|
||||
target_status_or_phase=TaskStatus.CANCELLED,
|
||||
reason="office_recovery_cancel",
|
||||
release_claim=True,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Recovery cancel skipped %s: %s", task.id, exc)
|
||||
failed_ids.append(projection_id)
|
||||
continue
|
||||
if task.status != TaskStatus.CANCELLED:
|
||||
logger.warning("Recovery cancel preserved non-cancelled phase for %s", task.id)
|
||||
failed_ids.append(projection_id)
|
||||
continue
|
||||
cancelled_count += 1
|
||||
|
||||
# Clean orphaned checkpoints
|
||||
await self._clean_orphaned_checkpoints(wf, tasks)
|
||||
|
||||
# Invalidate cache
|
||||
self._cache_until = 0.0
|
||||
|
||||
await self._set_run_recovery_state(
|
||||
wf.parent_session_id,
|
||||
status="cancelled",
|
||||
lifecycle_status="cancelled",
|
||||
)
|
||||
await self._broadcast_status()
|
||||
return {
|
||||
"ok": True,
|
||||
"cancelled_count": cancelled_count,
|
||||
"failed_work_item_projection_ids": failed_ids,
|
||||
}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _execute_recovery(
|
||||
self,
|
||||
parent_task_id: str,
|
||||
parent_session_id: str,
|
||||
plan: Any,
|
||||
tasks: list[Any],
|
||||
) -> None:
|
||||
"""Run the work-item executor in the background."""
|
||||
project_id = self._engine.project_id or "default"
|
||||
try:
|
||||
await self._set_run_recovery_state(
|
||||
parent_session_id,
|
||||
status="started",
|
||||
lifecycle_status="active",
|
||||
)
|
||||
await self._broadcast({"type": "recovery_result", "payload": {
|
||||
"project_id": project_id,
|
||||
"parent_task_id": parent_task_id, "status": "started",
|
||||
}})
|
||||
|
||||
executor = self._engine.company_executor
|
||||
if not executor:
|
||||
raise RuntimeError("company_executor not available")
|
||||
|
||||
result = await executor.execute(plan, tasks)
|
||||
|
||||
await self._broadcast({"type": "recovery_result", "payload": {
|
||||
"project_id": project_id,
|
||||
"parent_task_id": parent_task_id, "status": "completed",
|
||||
"summary": result[:500] if result else "",
|
||||
}})
|
||||
await self._set_run_recovery_state(
|
||||
parent_session_id,
|
||||
status="completed",
|
||||
lifecycle_status="active",
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
await self._broadcast({"type": "recovery_result", "payload": {
|
||||
"project_id": project_id,
|
||||
"parent_task_id": parent_task_id, "status": "cancelled",
|
||||
}})
|
||||
await self._set_run_recovery_state(
|
||||
parent_session_id,
|
||||
status="cancelled",
|
||||
lifecycle_status="cancelled",
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(f"Recovery execution failed for {parent_task_id}: {exc}")
|
||||
await self._broadcast({"type": "recovery_result", "payload": {
|
||||
"project_id": project_id,
|
||||
"parent_task_id": parent_task_id, "status": "failed",
|
||||
"error": str(exc),
|
||||
}})
|
||||
await self._set_run_recovery_state(
|
||||
parent_session_id,
|
||||
status="failed",
|
||||
lifecycle_status="blocked",
|
||||
extra={"error": str(exc)},
|
||||
)
|
||||
finally:
|
||||
self._active_recoveries.pop(parent_task_id, None)
|
||||
self._cache_until = 0.0
|
||||
await self._broadcast_status()
|
||||
|
||||
async def _clean_orphaned_checkpoints(
|
||||
self,
|
||||
wf: InterruptedWorkItemRun,
|
||||
tasks: list[Any],
|
||||
) -> int:
|
||||
"""Resolve pending checkpoints whose tasks are no longer active."""
|
||||
store = self._engine.store
|
||||
if not store:
|
||||
return 0
|
||||
|
||||
session_ids = {
|
||||
str(getattr(t, "session_id", "") or "").strip()
|
||||
for t in tasks
|
||||
}
|
||||
session_ids.add(wf.parent_session_id)
|
||||
session_ids.discard("")
|
||||
|
||||
cleaned = 0
|
||||
try:
|
||||
pending = await store.get_pending_checkpoints(
|
||||
project_id=wf.project_id,
|
||||
)
|
||||
for cp in pending:
|
||||
cp_session = str(cp.session_id or "").strip()
|
||||
if cp_session in session_ids:
|
||||
await store.resolve_execution_checkpoint(
|
||||
cp.checkpoint_id, status="cancelled"
|
||||
)
|
||||
cleaned += 1
|
||||
except Exception as exc:
|
||||
logger.debug(f"Checkpoint cleanup error: {exc}")
|
||||
|
||||
return cleaned
|
||||
|
||||
async def _broadcast_status(self) -> None:
|
||||
"""Push updated recovery status to all connected clients."""
|
||||
try:
|
||||
status = await self.get_recovery_status()
|
||||
await self._broadcast({"type": "recovery_status", "payload":
|
||||
_serialize_status(status, project_id=self._engine.project_id or "default")
|
||||
})
|
||||
except Exception as exc:
|
||||
logger.debug(f"Recovery status broadcast failed: {exc}")
|
||||
|
||||
def _invalidate_cache(self) -> None:
|
||||
"""Force next get_recovery_status to re-scan."""
|
||||
self._cache_until = 0.0
|
||||
self._cached = None
|
||||
|
||||
async def _set_run_recovery_state(
|
||||
self,
|
||||
key: str,
|
||||
*,
|
||||
status: str,
|
||||
lifecycle_status: str | None = None,
|
||||
match_task_id: bool = False,
|
||||
extra: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
store = getattr(self._engine, "store", None)
|
||||
if not store or not hasattr(store, "list_delegation_runs") or not hasattr(store, "save_delegation_run"):
|
||||
return
|
||||
runs = await store.list_delegation_runs(project_id=self._engine.project_id or "default")
|
||||
target = None
|
||||
if match_task_id:
|
||||
for run in runs:
|
||||
metadata = dict(getattr(run, "metadata", {}) or {})
|
||||
if str(metadata.get("origin_task_id", "") or "").strip() == key:
|
||||
target = run
|
||||
break
|
||||
if target is None:
|
||||
for run in runs:
|
||||
if str(run.session_id or "").strip() == key:
|
||||
target = run
|
||||
break
|
||||
if target is None:
|
||||
return
|
||||
target.recovery_pointer = {
|
||||
**dict(getattr(target, "recovery_pointer", {}) or {}),
|
||||
"status": status,
|
||||
"updated_at": datetime.now().isoformat(),
|
||||
**dict(extra or {}),
|
||||
}
|
||||
if lifecycle_status:
|
||||
target.lifecycle_status = lifecycle_status
|
||||
target.updated_at = datetime.now()
|
||||
await store.save_delegation_run(target)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _is_interrupted(task: Any) -> bool:
|
||||
"""Check if a task was interrupted by crash."""
|
||||
from opc.core.models import TaskStatus
|
||||
if task.status != TaskStatus.FAILED:
|
||||
return False
|
||||
meta = getattr(task, "metadata", {}) or {}
|
||||
if meta.get("interrupted_recovery"):
|
||||
return True
|
||||
result = getattr(task, "result", {}) or {}
|
||||
artifacts = result.get("artifacts", {}) or {}
|
||||
return bool(artifacts.get("interrupted"))
|
||||
|
||||
|
||||
def _serialize_status(status: RecoveryStatus, *, project_id: str | None = None) -> dict[str, Any]:
|
||||
"""Convert RecoveryStatus to JSON-safe dict."""
|
||||
resolved_project_id = str(project_id or "").strip()
|
||||
if not resolved_project_id:
|
||||
for item in status.interrupted:
|
||||
if item.project_id:
|
||||
resolved_project_id = item.project_id
|
||||
break
|
||||
payload = {
|
||||
"interrupted": [
|
||||
{
|
||||
"parent_session_id": w.parent_session_id,
|
||||
"parent_task_id": w.parent_task_id,
|
||||
"project_id": w.project_id,
|
||||
"title": w.title,
|
||||
"profile": w.profile,
|
||||
"interrupted_at": w.interrupted_at,
|
||||
"work_items": [
|
||||
{
|
||||
"work_item_projection_id": s.work_item_projection_id,
|
||||
"title": s.title,
|
||||
"task_id": s.task_id,
|
||||
"status": s.status,
|
||||
"interrupted": s.interrupted,
|
||||
"previous_status": s.previous_status,
|
||||
}
|
||||
for s in w.work_items
|
||||
],
|
||||
}
|
||||
for w in status.interrupted
|
||||
],
|
||||
"active_recoveries": status.active_recoveries,
|
||||
"scanned_at": status.scanned_at,
|
||||
}
|
||||
if resolved_project_id:
|
||||
payload["project_id"] = resolved_project_id
|
||||
return payload
|
||||
Reference in New Issue
Block a user