314 lines
11 KiB
Python
314 lines
11 KiB
Python
"""Provider-neutral worker runtime for AIOS external executors.
|
|
|
|
Providers implement only step execution and notification delivery. Discovery,
|
|
dependency gates, lifecycle truth, heartbeat, evidence validation and durable
|
|
notification acknowledgement are identical for every current and future
|
|
executor Provider.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Mapping, Protocol, Sequence
|
|
|
|
from .execution_plan import ExecutionPlanCoordinator, ExecutionStageManifest
|
|
from .executor_protocol import (
|
|
ExecutorIdentity,
|
|
ExecutorLifecycleSession,
|
|
ExecutorProviderRegistry,
|
|
normalize_executor,
|
|
)
|
|
from .external_task_control import ExternalTaskCoordinator, NotificationRecord
|
|
from .task_execution_registry import ExecutionTruthState
|
|
|
|
|
|
class StepRunner(Protocol):
|
|
def run_step(
|
|
self,
|
|
*,
|
|
plan_id: str,
|
|
stage: ExecutionStageManifest,
|
|
step: str,
|
|
step_number: int,
|
|
executor_id: str,
|
|
executor_run_id: str,
|
|
) -> Mapping[str, str]: ...
|
|
|
|
|
|
class NotificationSender(Protocol):
|
|
def send(self, notification: NotificationRecord) -> str: ...
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StageRunResult:
|
|
plan_id: str
|
|
stage_id: str
|
|
task_id: str
|
|
executor_type: str
|
|
executor_id: str
|
|
executor_run_id: str
|
|
state: str
|
|
evidence: Mapping[str, str] = field(default_factory=dict)
|
|
error: Mapping[str, Any] = field(default_factory=dict)
|
|
delivered_notifications: Sequence[str] = ()
|
|
|
|
def as_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"plan_id": self.plan_id,
|
|
"stage_id": self.stage_id,
|
|
"task_id": self.task_id,
|
|
"executor_type": self.executor_type,
|
|
"executor_id": self.executor_id,
|
|
"executor_run_id": self.executor_run_id,
|
|
"state": self.state,
|
|
"evidence": dict(self.evidence),
|
|
"error": dict(self.error),
|
|
"delivered_notifications": list(self.delivered_notifications),
|
|
}
|
|
|
|
|
|
class UnifiedExecutorWorker:
|
|
"""Dependency-aware worker shared by every registered Provider."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
plan_coordinator: ExecutionPlanCoordinator,
|
|
task_coordinator: ExternalTaskCoordinator,
|
|
providers: ExecutorProviderRegistry,
|
|
executor_type: str,
|
|
executor_id: str,
|
|
step_runner: StepRunner,
|
|
notification_sender: NotificationSender,
|
|
) -> None:
|
|
self.plan_coordinator = plan_coordinator
|
|
self.task_coordinator = task_coordinator
|
|
self.providers = providers
|
|
self.executor_type = normalize_executor(executor_type)
|
|
self.executor_id = str(executor_id or "").strip()
|
|
if not self.executor_id:
|
|
raise ValueError("executor_id is required")
|
|
self.providers.get(self.executor_type)
|
|
self.step_runner = step_runner
|
|
self.notification_sender = notification_sender
|
|
|
|
def _session(self, executor_run_id: str) -> ExecutorLifecycleSession:
|
|
return ExecutorLifecycleSession(
|
|
self.task_coordinator,
|
|
ExecutorIdentity(
|
|
executor_type=self.executor_type,
|
|
executor_id=self.executor_id,
|
|
executor_run_id=executor_run_id,
|
|
),
|
|
registry=self.providers,
|
|
)
|
|
|
|
def _notification_candidates(self, task_id: str) -> tuple[NotificationRecord, ...]:
|
|
"""Return idempotent delivery candidates, including prior failures.
|
|
|
|
The outbox has a unique ``(task_id, event_type)`` key. Retrying therefore
|
|
reuses the same notification ID and cannot create duplicate STARTED or
|
|
SUCCEEDED events. A later successful acknowledgement replaces FAILED
|
|
with DELIVERED and records the receipt hash.
|
|
"""
|
|
candidates: dict[str, NotificationRecord] = {}
|
|
for status in ("PENDING", "FAILED"):
|
|
for item in self.task_coordinator.notifications.list(
|
|
task_id=task_id,
|
|
status=status,
|
|
):
|
|
candidates[item.notification_id] = item
|
|
return tuple(
|
|
candidates[key]
|
|
for key in sorted(
|
|
candidates,
|
|
key=lambda item: (
|
|
candidates[item].created_at,
|
|
candidates[item].notification_id,
|
|
),
|
|
)
|
|
)
|
|
|
|
def _deliver_notifications(self, task_id: str) -> tuple[str, ...]:
|
|
delivered: list[str] = []
|
|
for item in self._notification_candidates(task_id):
|
|
try:
|
|
receipt = self.notification_sender.send(item)
|
|
acknowledged = self.task_coordinator.notifications.acknowledge(
|
|
item.notification_id,
|
|
receipt_sha256=receipt,
|
|
)
|
|
delivered.append(acknowledged.event_type)
|
|
except Exception as exc:
|
|
self.task_coordinator.notifications.fail(
|
|
item.notification_id,
|
|
error=f"{type(exc).__name__}: {exc}",
|
|
)
|
|
return tuple(sorted(set(delivered)))
|
|
|
|
def retry_failed_notifications(self, *, task_id: str) -> tuple[str, ...]:
|
|
"""Retry failed owner notifications without changing task truth."""
|
|
return self._deliver_notifications(task_id)
|
|
|
|
def run_ready(
|
|
self,
|
|
*,
|
|
plan_id: str,
|
|
max_stages: int = 1,
|
|
) -> list[StageRunResult]:
|
|
if max_stages < 1:
|
|
raise ValueError("max_stages must be positive")
|
|
ready = self.plan_coordinator.discover_ready(
|
|
plan_id=plan_id,
|
|
executor_type=self.executor_type,
|
|
)[:max_stages]
|
|
return [
|
|
self.run_stage(
|
|
plan_id=plan_id,
|
|
stage_id=item["stage"]["stage_id"],
|
|
)
|
|
for item in ready
|
|
]
|
|
|
|
def run_stage(
|
|
self,
|
|
*,
|
|
plan_id: str,
|
|
stage_id: str,
|
|
executor_run_id: str | None = None,
|
|
) -> StageRunResult:
|
|
plan = self.plan_coordinator.plan_store.get(plan_id)
|
|
stage = plan.stage(stage_id)
|
|
if stage.executor_type != self.executor_type:
|
|
raise PermissionError(
|
|
f"stage executor is {stage.executor_type}, not {self.executor_type}"
|
|
)
|
|
dependency = self.plan_coordinator.dependency_gate(plan_id, stage_id)
|
|
if not dependency["passed"]:
|
|
raise RuntimeError("stage dependencies are not satisfied")
|
|
|
|
run_id = executor_run_id or f"{self.executor_type}_run_{uuid.uuid4().hex}"
|
|
session = self._session(run_id)
|
|
record = self.task_coordinator.registry.get(stage.task_id)
|
|
evidence: dict[str, str] = dict(record.evidence)
|
|
delivered: list[str] = []
|
|
try:
|
|
if record.state == ExecutionTruthState.DECLARED:
|
|
record = session.register(stage.task_id)
|
|
if record.state == ExecutionTruthState.REGISTERED:
|
|
record = session.claim(stage.task_id)
|
|
if record.state == ExecutionTruthState.CLAIMED:
|
|
record = session.start(stage.task_id)
|
|
delivered.extend(self._deliver_notifications(stage.task_id))
|
|
if record.state != ExecutionTruthState.RUNNING:
|
|
raise RuntimeError(
|
|
f"stage cannot execute from state {record.state.value}"
|
|
)
|
|
|
|
# Re-running the current step is intentional: Provider steps must be
|
|
# idempotent and a crash between side effect and heartbeat cannot be
|
|
# safely interpreted as completed work.
|
|
start_step = max(1, record.current_step)
|
|
for number, step in enumerate(stage.steps, start=1):
|
|
if number < start_step:
|
|
continue
|
|
produced = self.step_runner.run_step(
|
|
plan_id=plan_id,
|
|
stage=stage,
|
|
step=step,
|
|
step_number=number,
|
|
executor_id=self.executor_id,
|
|
executor_run_id=run_id,
|
|
)
|
|
if not isinstance(produced, Mapping):
|
|
raise TypeError("step runner must return an evidence mapping")
|
|
evidence.update(
|
|
{
|
|
str(key).strip(): str(value).strip().lower()
|
|
for key, value in produced.items()
|
|
}
|
|
)
|
|
record = session.heartbeat(stage.task_id, current_step=number)
|
|
|
|
record = session.succeed(stage.task_id, evidence=evidence)
|
|
delivered.extend(self._deliver_notifications(stage.task_id))
|
|
return StageRunResult(
|
|
plan_id=plan_id,
|
|
stage_id=stage.stage_id,
|
|
task_id=stage.task_id,
|
|
executor_type=self.executor_type,
|
|
executor_id=self.executor_id,
|
|
executor_run_id=run_id,
|
|
state=record.state.value,
|
|
evidence=dict(record.evidence),
|
|
delivered_notifications=tuple(sorted(set(delivered))),
|
|
)
|
|
except Exception as exc:
|
|
current = self.task_coordinator.registry.get(stage.task_id)
|
|
if current.state in {
|
|
ExecutionTruthState.REGISTERED,
|
|
ExecutionTruthState.CLAIMED,
|
|
ExecutionTruthState.RUNNING,
|
|
}:
|
|
current = session.fail(
|
|
stage.task_id,
|
|
error={
|
|
"code": "executor_stage_failed",
|
|
"error_type": type(exc).__name__,
|
|
"message": str(exc),
|
|
},
|
|
evidence=evidence or None,
|
|
)
|
|
delivered.extend(self._deliver_notifications(stage.task_id))
|
|
return StageRunResult(
|
|
plan_id=plan_id,
|
|
stage_id=stage.stage_id,
|
|
task_id=stage.task_id,
|
|
executor_type=self.executor_type,
|
|
executor_id=self.executor_id,
|
|
executor_run_id=run_id,
|
|
state=current.state.value,
|
|
evidence=dict(current.evidence),
|
|
error=dict(current.error),
|
|
delivered_notifications=tuple(sorted(set(delivered))),
|
|
)
|
|
|
|
|
|
class JsonEvidenceStepRunner:
|
|
"""Adapter for handlers returning evidence-name to SHA256 mappings."""
|
|
|
|
def __init__(self, handler) -> None:
|
|
self.handler = handler
|
|
|
|
def run_step(self, **values: Any) -> Mapping[str, str]:
|
|
result = self.handler(**values)
|
|
if not isinstance(result, Mapping):
|
|
raise TypeError("step handler must return a mapping")
|
|
return {str(key): str(value) for key, value in result.items()}
|
|
|
|
|
|
class JsonReceiptNotificationSender:
|
|
"""Adapter for senders returning a SHA256 or receipt JSON object."""
|
|
|
|
def __init__(self, sender) -> None:
|
|
self.sender = sender
|
|
|
|
def send(self, notification: NotificationRecord) -> str:
|
|
result = self.sender(notification.as_dict())
|
|
if isinstance(result, str):
|
|
return result
|
|
if not isinstance(result, Mapping):
|
|
raise TypeError(
|
|
"notification sender must return SHA256 or a receipt object"
|
|
)
|
|
encoded = json.dumps(
|
|
dict(result),
|
|
ensure_ascii=False,
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
default=str,
|
|
).encode("utf-8")
|
|
return hashlib.sha256(encoded).hexdigest()
|