LZH-YS1998
553 lines
21 KiB
Python
553 lines
21 KiB
Python
"""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
|