265 lines
9.0 KiB
Python
265 lines
9.0 KiB
Python
"""AIOS V2 execution orchestrator.
|
|
|
|
The orchestrator is the only component allowed to translate an approved AIOS
|
|
task into an OpenClaw command or Hermes job. OpenClaw and Hermes remain pure
|
|
execution runtimes and never own tenant policy or task lifecycle.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any, Mapping
|
|
|
|
from aios_hermes import HermesAdapter, HermesJob, HermesResult
|
|
from aios_openclaw import (
|
|
CommandState,
|
|
OpenClawAdapter,
|
|
OpenClawCommand,
|
|
OpenClawResult,
|
|
sign_command,
|
|
)
|
|
|
|
from .control_plane import (
|
|
AgentRecord,
|
|
ControlPlaneService,
|
|
RequestContext,
|
|
TaskRecord,
|
|
TaskState,
|
|
)
|
|
|
|
|
|
class OrchestrationError(RuntimeError):
|
|
"""AIOS failed to dispatch or reconcile a runtime execution."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DispatchRecord:
|
|
tenant_id: str
|
|
task_id: str
|
|
agent_id: str
|
|
runtime: str
|
|
runtime_execution_id: str
|
|
state: str
|
|
trace_id: str
|
|
|
|
|
|
class ExecutionOrchestrator:
|
|
"""Dispatch approved AIOS tasks to runtime adapters and reconcile results."""
|
|
|
|
def __init__(
|
|
self,
|
|
control_plane: ControlPlaneService,
|
|
*,
|
|
openclaw: OpenClawAdapter | None = None,
|
|
hermes: HermesAdapter | None = None,
|
|
openclaw_signing_material_resolver=None,
|
|
openclaw_key_id: str = "default",
|
|
default_cell_id: str = "default",
|
|
default_profile_id: str = "default",
|
|
) -> None:
|
|
self.control_plane = control_plane
|
|
self.openclaw = openclaw
|
|
self.hermes = hermes
|
|
self.openclaw_signing_material_resolver = openclaw_signing_material_resolver
|
|
self.openclaw_key_id = openclaw_key_id
|
|
self.default_cell_id = default_cell_id
|
|
self.default_profile_id = default_profile_id
|
|
self._dispatches: dict[tuple[str, str], DispatchRecord] = {}
|
|
|
|
def dispatch(
|
|
self,
|
|
context: RequestContext,
|
|
*,
|
|
task_id: str,
|
|
timeout_seconds: float | None = None,
|
|
) -> DispatchRecord:
|
|
task = self.control_plane.get_task(context, task_id)
|
|
if task.state != TaskState.APPROVED:
|
|
raise OrchestrationError("only APPROVED tasks may be dispatched")
|
|
agent = self.control_plane.get_agent(context, task.agent_id)
|
|
existing = self._dispatches.get((task.tenant_id, task.task_id))
|
|
if existing is not None:
|
|
return existing
|
|
|
|
running = self.control_plane.transition_task(
|
|
context,
|
|
task_id=task.task_id,
|
|
target_state=TaskState.RUNNING,
|
|
expected_revision=task.revision,
|
|
)
|
|
|
|
runtime = agent.runtime.strip().lower()
|
|
try:
|
|
if runtime == "openclaw":
|
|
record = self._dispatch_openclaw(
|
|
context,
|
|
agent,
|
|
running,
|
|
timeout_seconds=timeout_seconds,
|
|
)
|
|
elif runtime == "hermes":
|
|
record = self._dispatch_hermes(context, agent, running)
|
|
else:
|
|
raise OrchestrationError(f"unsupported runtime: {agent.runtime}")
|
|
except Exception as exc:
|
|
current = self.control_plane.get_task(context, task.task_id)
|
|
if current.state == TaskState.RUNNING:
|
|
self.control_plane.transition_task(
|
|
context,
|
|
task_id=task.task_id,
|
|
target_state=TaskState.FAILED,
|
|
expected_revision=current.revision,
|
|
error={"code": "dispatch_failed", "message": str(exc)},
|
|
)
|
|
raise
|
|
|
|
self._dispatches[(task.tenant_id, task.task_id)] = record
|
|
return record
|
|
|
|
def _dispatch_openclaw(
|
|
self,
|
|
context: RequestContext,
|
|
agent: AgentRecord,
|
|
task: TaskRecord,
|
|
*,
|
|
timeout_seconds: float | None,
|
|
) -> DispatchRecord:
|
|
if self.openclaw is None or self.openclaw_signing_material_resolver is None:
|
|
raise OrchestrationError("OpenClaw runtime is not configured")
|
|
key_material = self.openclaw_signing_material_resolver(
|
|
task.tenant_id,
|
|
self.openclaw_key_id,
|
|
)
|
|
command = OpenClawCommand.create(
|
|
tenant_id=task.tenant_id,
|
|
task_id=task.task_id,
|
|
trace_id=context.trace_id,
|
|
action=task.action,
|
|
payload=task.payload,
|
|
idempotency_key=task.idempotency_key or task.task_id,
|
|
key_id=self.openclaw_key_id,
|
|
)
|
|
command = sign_command(command, key_material)
|
|
result = self.openclaw.execute(command, timeout_seconds=timeout_seconds)
|
|
self.reconcile_openclaw(context, result)
|
|
return DispatchRecord(
|
|
tenant_id=task.tenant_id,
|
|
task_id=task.task_id,
|
|
agent_id=agent.agent_id,
|
|
runtime="openclaw",
|
|
runtime_execution_id=command.command_id,
|
|
state=result.state.value,
|
|
trace_id=context.trace_id,
|
|
)
|
|
|
|
def _dispatch_hermes(
|
|
self,
|
|
context: RequestContext,
|
|
agent: AgentRecord,
|
|
task: TaskRecord,
|
|
) -> DispatchRecord:
|
|
if self.hermes is None:
|
|
raise OrchestrationError("Hermes runtime is not configured")
|
|
metadata = dict(agent.metadata)
|
|
job = HermesJob.create(
|
|
tenant_id=task.tenant_id,
|
|
cell_id=str(metadata.get("cell_id") or self.default_cell_id),
|
|
profile_id=str(metadata.get("profile_id") or self.default_profile_id),
|
|
task_id=task.task_id,
|
|
trace_id=context.trace_id,
|
|
action=task.action,
|
|
payload=task.payload,
|
|
idempotency_key=task.idempotency_key or task.task_id,
|
|
priority=int(metadata.get("priority", 100)),
|
|
)
|
|
queued = self.hermes.enqueue(job)
|
|
return DispatchRecord(
|
|
tenant_id=task.tenant_id,
|
|
task_id=task.task_id,
|
|
agent_id=agent.agent_id,
|
|
runtime="hermes",
|
|
runtime_execution_id=queued.job_id,
|
|
state=queued.state.value,
|
|
trace_id=context.trace_id,
|
|
)
|
|
|
|
def reconcile_openclaw(
|
|
self,
|
|
context: RequestContext,
|
|
result: OpenClawResult,
|
|
) -> TaskRecord:
|
|
task = self.control_plane.get_task(context, result.task_id)
|
|
if result.tenant_id != task.tenant_id:
|
|
raise OrchestrationError("OpenClaw result tenant mismatch")
|
|
if task.state != TaskState.RUNNING:
|
|
return task
|
|
if result.state == CommandState.COMPLETED:
|
|
return self.control_plane.transition_task(
|
|
context,
|
|
task_id=task.task_id,
|
|
target_state=TaskState.COMPLETED,
|
|
expected_revision=task.revision,
|
|
result=result.output or {},
|
|
)
|
|
if result.state in {
|
|
CommandState.FAILED,
|
|
CommandState.TIMED_OUT,
|
|
CommandState.REJECTED,
|
|
}:
|
|
return self.control_plane.transition_task(
|
|
context,
|
|
task_id=task.task_id,
|
|
target_state=TaskState.FAILED,
|
|
expected_revision=task.revision,
|
|
error=result.error or {"code": result.state.value.lower()},
|
|
)
|
|
if result.state == CommandState.CANCELLED:
|
|
return self.control_plane.transition_task(
|
|
context,
|
|
task_id=task.task_id,
|
|
target_state=TaskState.CANCELLED,
|
|
expected_revision=task.revision,
|
|
)
|
|
return task
|
|
|
|
def reconcile_hermes(
|
|
self,
|
|
context: RequestContext,
|
|
result: HermesResult,
|
|
) -> TaskRecord:
|
|
task = self.control_plane.get_task(context, result.task_id)
|
|
if result.tenant_id != task.tenant_id:
|
|
raise OrchestrationError("Hermes result tenant mismatch")
|
|
if task.state != TaskState.RUNNING:
|
|
return task
|
|
if result.state.value == "COMPLETED":
|
|
return self.control_plane.transition_task(
|
|
context,
|
|
task_id=task.task_id,
|
|
target_state=TaskState.COMPLETED,
|
|
expected_revision=task.revision,
|
|
result=result.output or {},
|
|
)
|
|
if result.state.value == "FAILED":
|
|
return self.control_plane.transition_task(
|
|
context,
|
|
task_id=task.task_id,
|
|
target_state=TaskState.FAILED,
|
|
expected_revision=task.revision,
|
|
error=result.error or {"code": "hermes_failed"},
|
|
)
|
|
if result.state.value == "CANCELLED":
|
|
return self.control_plane.transition_task(
|
|
context,
|
|
task_id=task.task_id,
|
|
target_state=TaskState.CANCELLED,
|
|
expected_revision=task.revision,
|
|
)
|
|
return task
|
|
|
|
def dispatch_snapshot(self, *, tenant_id: str) -> Mapping[str, DispatchRecord]:
|
|
return {
|
|
task_id: record
|
|
for (record_tenant, task_id), record in self._dispatches.items()
|
|
if record_tenant == tenant_id
|
|
}
|