600 lines
22 KiB
Python
600 lines
22 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. Every dispatch is persisted before
|
|
Runtime side effects. Runtime results must bind to the exact dispatch,
|
|
command/job, tenant, task, trace, principal and provenance context.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Mapping, Protocol
|
|
|
|
from aios_hermes import HermesAdapter, HermesJob, HermesJobState, HermesResult
|
|
from aios_openclaw import (
|
|
CommandState,
|
|
OpenClawAdapter,
|
|
OpenClawCommand,
|
|
OpenClawResult,
|
|
sign_command,
|
|
)
|
|
|
|
from .control_plane import (
|
|
AgentRecord,
|
|
ControlPlaneService,
|
|
RequestContext,
|
|
TaskRecord,
|
|
TaskState,
|
|
)
|
|
from .dispatch_store import (
|
|
DispatchConflict,
|
|
DispatchRecord,
|
|
DispatchRepository,
|
|
DispatchState,
|
|
InMemoryDispatchRepository,
|
|
TERMINAL_DISPATCH_STATES,
|
|
)
|
|
from .provenance import ProvenanceDenied, digest_payload
|
|
|
|
|
|
class RuntimeAuthorizationWriter(Protocol):
|
|
def authorize_runtime_execute(self, **values: Any) -> Any: ...
|
|
|
|
|
|
class OrchestrationError(RuntimeError):
|
|
"""AIOS failed to dispatch or reconcile a runtime execution."""
|
|
|
|
|
|
class ExecutionOrchestrator:
|
|
"""Dispatch approved AIOS tasks and strictly reconcile Runtime 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",
|
|
provenance_writer: RuntimeAuthorizationWriter | None = None,
|
|
require_provenance: bool = False,
|
|
dispatch_repository: DispatchRepository | None = None,
|
|
) -> 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 = str(openclaw_key_id or "").strip() or "default"
|
|
self.default_cell_id = default_cell_id
|
|
self.default_profile_id = default_profile_id
|
|
self.provenance_writer = provenance_writer
|
|
self.require_provenance = bool(require_provenance)
|
|
self.dispatch_repository = dispatch_repository or InMemoryDispatchRepository()
|
|
if self.require_provenance and self.provenance_writer is None:
|
|
raise ValueError("production orchestration requires a provenance writer")
|
|
|
|
@staticmethod
|
|
def _record_provenance(value: str | None, *, required: bool, label: str) -> str:
|
|
normalized = str(value or "").strip()
|
|
if required and not normalized:
|
|
raise OrchestrationError(f"production dispatch requires {label}")
|
|
return normalized or f"{label}-not-required"
|
|
|
|
def _authorize_dispatch(
|
|
self,
|
|
context: RequestContext,
|
|
*,
|
|
task: TaskRecord,
|
|
agent: AgentRecord,
|
|
provenance_id: str | None,
|
|
run_id: str | None,
|
|
timeout_seconds: float | None,
|
|
) -> str:
|
|
if not self.require_provenance:
|
|
return "authorization-not-required"
|
|
if self.provenance_writer is None:
|
|
raise OrchestrationError("Run Provenance writer is not configured")
|
|
if not provenance_id or not run_id:
|
|
raise OrchestrationError(
|
|
"production dispatch requires provenance_id and run_id"
|
|
)
|
|
request_digest = digest_payload(
|
|
{
|
|
"operation": "dispatch",
|
|
"tenant_id": task.tenant_id,
|
|
"task_id": task.task_id,
|
|
"agent_id": agent.agent_id,
|
|
"principal_id": context.principal_id,
|
|
"runtime": agent.runtime.strip().lower(),
|
|
"action": task.action,
|
|
"payload": task.payload,
|
|
"trace_id": context.trace_id,
|
|
"timeout_seconds": timeout_seconds,
|
|
"task_revision": task.revision,
|
|
}
|
|
)
|
|
try:
|
|
authorization = self.provenance_writer.authorize_runtime_execute(
|
|
provenance_id=provenance_id,
|
|
run_id=run_id,
|
|
tenant_id=task.tenant_id,
|
|
task_id=task.task_id,
|
|
agent_id=agent.agent_id,
|
|
principal_id=context.principal_id,
|
|
request_digest=request_digest,
|
|
)
|
|
return authorization.authorization_id
|
|
except ProvenanceDenied as exc:
|
|
raise OrchestrationError(f"Run Provenance denied dispatch: {exc}") from exc
|
|
|
|
def dispatch(
|
|
self,
|
|
context: RequestContext,
|
|
*,
|
|
task_id: str,
|
|
timeout_seconds: float | None = None,
|
|
provenance_id: str | None = None,
|
|
run_id: str | None = None,
|
|
) -> DispatchRecord:
|
|
task = self.control_plane.get_task(context, task_id)
|
|
agent = self.control_plane.get_agent(context, task.agent_id)
|
|
persisted_provenance = self._record_provenance(
|
|
provenance_id, required=self.require_provenance, label="provenance_id"
|
|
)
|
|
persisted_run = self._record_provenance(
|
|
run_id, required=self.require_provenance, label="run_id"
|
|
)
|
|
|
|
existing = self.dispatch_repository.get_by_task(
|
|
task.task_id, tenant_id=task.tenant_id
|
|
)
|
|
if existing is not None:
|
|
if (
|
|
existing.provenance_id != persisted_provenance
|
|
or existing.run_id != persisted_run
|
|
or existing.principal_id != context.principal_id
|
|
):
|
|
raise OrchestrationError(
|
|
"existing dispatch belongs to another provenance or principal context"
|
|
)
|
|
return existing
|
|
|
|
if task.state != TaskState.APPROVED:
|
|
raise OrchestrationError("only APPROVED tasks may be dispatched")
|
|
|
|
authorization_id = self._authorize_dispatch(
|
|
context,
|
|
task=task,
|
|
agent=agent,
|
|
provenance_id=provenance_id,
|
|
run_id=run_id,
|
|
timeout_seconds=timeout_seconds,
|
|
)
|
|
runtime = agent.runtime.strip().lower()
|
|
prepared: DispatchRecord | None = None
|
|
|
|
try:
|
|
if runtime == "openclaw":
|
|
command = self._create_openclaw_command(context, task)
|
|
prepared = self.dispatch_repository.create(
|
|
DispatchRecord.create(
|
|
tenant_id=task.tenant_id,
|
|
task_id=task.task_id,
|
|
agent_id=agent.agent_id,
|
|
principal_id=context.principal_id,
|
|
runtime="openclaw",
|
|
runtime_execution_id=command.command_id,
|
|
trace_id=context.trace_id,
|
|
provenance_id=persisted_provenance,
|
|
run_id=persisted_run,
|
|
authorization_id=authorization_id,
|
|
)
|
|
)
|
|
self.control_plane.transition_task(
|
|
context,
|
|
task_id=task.task_id,
|
|
target_state=TaskState.RUNNING,
|
|
expected_revision=task.revision,
|
|
)
|
|
self.dispatch_repository.update(
|
|
prepared.dispatch_id,
|
|
tenant_id=prepared.tenant_id,
|
|
state=DispatchState.RUNNING,
|
|
)
|
|
if self.openclaw is None:
|
|
raise OrchestrationError("OpenClaw runtime is not configured")
|
|
result = self.openclaw.execute(
|
|
command, timeout_seconds=timeout_seconds
|
|
)
|
|
self.reconcile_openclaw(
|
|
context,
|
|
result,
|
|
dispatch_id=prepared.dispatch_id,
|
|
provenance_id=persisted_provenance,
|
|
run_id=persisted_run,
|
|
)
|
|
return self.dispatch_repository.get(
|
|
prepared.dispatch_id, tenant_id=prepared.tenant_id
|
|
)
|
|
|
|
if runtime == "hermes":
|
|
if self.hermes is None:
|
|
raise OrchestrationError("Hermes runtime is not configured")
|
|
job = self._create_hermes_job(context, agent, task)
|
|
prepared = self.dispatch_repository.create(
|
|
DispatchRecord.create(
|
|
tenant_id=task.tenant_id,
|
|
task_id=task.task_id,
|
|
agent_id=agent.agent_id,
|
|
principal_id=context.principal_id,
|
|
runtime="hermes",
|
|
runtime_execution_id=job.job_id,
|
|
trace_id=context.trace_id,
|
|
provenance_id=persisted_provenance,
|
|
run_id=persisted_run,
|
|
authorization_id=authorization_id,
|
|
)
|
|
)
|
|
self.control_plane.transition_task(
|
|
context,
|
|
task_id=task.task_id,
|
|
target_state=TaskState.RUNNING,
|
|
expected_revision=task.revision,
|
|
)
|
|
queued = self.hermes.enqueue(job)
|
|
return self.dispatch_repository.update(
|
|
prepared.dispatch_id,
|
|
tenant_id=prepared.tenant_id,
|
|
state=DispatchState.QUEUED,
|
|
result_digest=digest_payload(
|
|
{
|
|
"queued_job_id": queued.job_id,
|
|
"state": queued.state.value,
|
|
"revision": queued.revision,
|
|
}
|
|
),
|
|
)
|
|
|
|
raise OrchestrationError(f"unsupported runtime: {agent.runtime}")
|
|
except Exception as exc:
|
|
if prepared is not None:
|
|
try:
|
|
self.dispatch_repository.update(
|
|
prepared.dispatch_id,
|
|
tenant_id=prepared.tenant_id,
|
|
state=DispatchState.FAILED,
|
|
result_digest=digest_payload(
|
|
{
|
|
"code": "dispatch_failed",
|
|
"error_type": type(exc).__name__,
|
|
"message": str(exc),
|
|
}
|
|
),
|
|
)
|
|
except Exception:
|
|
pass
|
|
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
|
|
|
|
def _create_openclaw_command(
|
|
self,
|
|
context: RequestContext,
|
|
task: TaskRecord,
|
|
) -> OpenClawCommand:
|
|
if self.openclaw_signing_material_resolver is None:
|
|
raise OrchestrationError("OpenClaw signing material 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,
|
|
)
|
|
return sign_command(command, key_material)
|
|
|
|
def _create_hermes_job(
|
|
self,
|
|
context: RequestContext,
|
|
agent: AgentRecord,
|
|
task: TaskRecord,
|
|
) -> HermesJob:
|
|
metadata = dict(agent.metadata)
|
|
return 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)),
|
|
)
|
|
|
|
@staticmethod
|
|
def _assert_dispatch_context(
|
|
dispatch: DispatchRecord,
|
|
*,
|
|
runtime: str,
|
|
runtime_execution_id: str,
|
|
tenant_id: str,
|
|
task_id: str,
|
|
trace_id: str,
|
|
principal_id: str,
|
|
provenance_id: str,
|
|
run_id: str,
|
|
) -> None:
|
|
expected = {
|
|
"runtime": runtime,
|
|
"runtime_execution_id": runtime_execution_id,
|
|
"tenant_id": tenant_id,
|
|
"task_id": task_id,
|
|
"trace_id": trace_id,
|
|
"principal_id": principal_id,
|
|
"provenance_id": provenance_id,
|
|
"run_id": run_id,
|
|
}
|
|
for field_name, expected_value in expected.items():
|
|
if getattr(dispatch, field_name) != expected_value:
|
|
raise OrchestrationError(f"dispatch {field_name} mismatch")
|
|
|
|
def hermes_lifecycle_context(
|
|
self,
|
|
context: RequestContext,
|
|
*,
|
|
dispatch_id: str,
|
|
job_id: str,
|
|
tenant_id: str,
|
|
task_id: str,
|
|
trace_id: str,
|
|
principal_id: str,
|
|
provenance_id: str,
|
|
run_id: str,
|
|
) -> tuple[DispatchRecord, HermesJob]:
|
|
"""Resolve a runtime lifecycle request against its durable dispatch."""
|
|
dispatch = self.dispatch_repository.get(
|
|
dispatch_id, tenant_id=context.tenant_id
|
|
)
|
|
self._assert_dispatch_context(
|
|
dispatch,
|
|
runtime="hermes",
|
|
runtime_execution_id=job_id,
|
|
tenant_id=tenant_id,
|
|
task_id=task_id,
|
|
trace_id=trace_id,
|
|
principal_id=principal_id,
|
|
provenance_id=provenance_id,
|
|
run_id=run_id,
|
|
)
|
|
if context.tenant_id != tenant_id:
|
|
raise OrchestrationError("runtime context tenant_id mismatch")
|
|
if context.principal_id != principal_id:
|
|
raise OrchestrationError("runtime context principal_id mismatch")
|
|
if context.trace_id != trace_id:
|
|
raise OrchestrationError("runtime context trace_id mismatch")
|
|
if self.hermes is None:
|
|
raise OrchestrationError("Hermes runtime is not configured")
|
|
job = self.hermes.get(job_id, tenant_id=tenant_id)
|
|
if job.task_id != task_id or job.trace_id != trace_id:
|
|
raise OrchestrationError("Hermes job identity mismatch")
|
|
return dispatch, job
|
|
|
|
def bind_hermes_worker(
|
|
self,
|
|
context: RequestContext,
|
|
job: HermesJob,
|
|
*,
|
|
dispatch_id: str,
|
|
provenance_id: str,
|
|
run_id: str,
|
|
worker_id: str,
|
|
) -> DispatchRecord:
|
|
dispatch, current = self.hermes_lifecycle_context(
|
|
context,
|
|
dispatch_id=dispatch_id,
|
|
job_id=job.job_id,
|
|
tenant_id=job.tenant_id,
|
|
task_id=job.task_id,
|
|
trace_id=job.trace_id,
|
|
principal_id=context.principal_id,
|
|
provenance_id=provenance_id,
|
|
run_id=run_id,
|
|
)
|
|
if current.assigned_worker_id != worker_id:
|
|
raise OrchestrationError("Hermes worker does not own claimed job")
|
|
return self.dispatch_repository.update(
|
|
dispatch.dispatch_id,
|
|
tenant_id=dispatch.tenant_id,
|
|
expected_worker_id=worker_id,
|
|
)
|
|
|
|
def reconcile_openclaw(
|
|
self,
|
|
context: RequestContext,
|
|
result: OpenClawResult,
|
|
*,
|
|
dispatch_id: str,
|
|
provenance_id: str,
|
|
run_id: str,
|
|
) -> TaskRecord:
|
|
dispatch = self.dispatch_repository.get(
|
|
dispatch_id, tenant_id=context.tenant_id
|
|
)
|
|
self._assert_dispatch_context(
|
|
dispatch,
|
|
runtime="openclaw",
|
|
runtime_execution_id=result.command_id,
|
|
tenant_id=result.tenant_id,
|
|
task_id=result.task_id,
|
|
trace_id=result.trace_id,
|
|
principal_id=context.principal_id,
|
|
provenance_id=provenance_id,
|
|
run_id=run_id,
|
|
)
|
|
result_digest = digest_payload(
|
|
{
|
|
"command_id": result.command_id,
|
|
"state": result.state.value,
|
|
"output": result.output,
|
|
"error": result.error,
|
|
"runtime_id": result.runtime_id,
|
|
"attempts": result.attempts,
|
|
"finished_at": result.finished_at,
|
|
}
|
|
)
|
|
if dispatch.state in TERMINAL_DISPATCH_STATES:
|
|
if dispatch.result_digest != result_digest:
|
|
raise DispatchConflict("terminal dispatch received a different result")
|
|
return self.control_plane.get_task(context, dispatch.task_id)
|
|
|
|
task = self.control_plane.get_task(context, dispatch.task_id)
|
|
if task.state != TaskState.RUNNING:
|
|
raise OrchestrationError("task is not RUNNING for active dispatch")
|
|
if result.state == CommandState.COMPLETED:
|
|
target = TaskState.COMPLETED
|
|
values = {"result": result.output or {}}
|
|
dispatch_state = DispatchState.COMPLETED
|
|
elif result.state in {
|
|
CommandState.FAILED,
|
|
CommandState.TIMED_OUT,
|
|
CommandState.REJECTED,
|
|
}:
|
|
target = TaskState.FAILED
|
|
values = {
|
|
"error": result.error or {"code": result.state.value.lower()}
|
|
}
|
|
dispatch_state = DispatchState.FAILED
|
|
elif result.state == CommandState.CANCELLED:
|
|
target = TaskState.CANCELLED
|
|
values = {}
|
|
dispatch_state = DispatchState.CANCELLED
|
|
else:
|
|
self.dispatch_repository.update(
|
|
dispatch.dispatch_id,
|
|
tenant_id=dispatch.tenant_id,
|
|
state=DispatchState.RUNNING,
|
|
result_digest=result_digest,
|
|
)
|
|
return task
|
|
|
|
updated_task = self.control_plane.transition_task(
|
|
context,
|
|
task_id=task.task_id,
|
|
target_state=target,
|
|
expected_revision=task.revision,
|
|
**values,
|
|
)
|
|
self.dispatch_repository.update(
|
|
dispatch.dispatch_id,
|
|
tenant_id=dispatch.tenant_id,
|
|
state=dispatch_state,
|
|
result_digest=result_digest,
|
|
)
|
|
return updated_task
|
|
|
|
def reconcile_hermes(
|
|
self,
|
|
context: RequestContext,
|
|
result: HermesResult,
|
|
*,
|
|
dispatch_id: str,
|
|
provenance_id: str,
|
|
run_id: str,
|
|
) -> TaskRecord:
|
|
dispatch = self.dispatch_repository.get(
|
|
dispatch_id, tenant_id=context.tenant_id
|
|
)
|
|
self._assert_dispatch_context(
|
|
dispatch,
|
|
runtime="hermes",
|
|
runtime_execution_id=result.job_id,
|
|
tenant_id=result.tenant_id,
|
|
task_id=result.task_id,
|
|
trace_id=result.trace_id,
|
|
principal_id=context.principal_id,
|
|
provenance_id=provenance_id,
|
|
run_id=run_id,
|
|
)
|
|
if self.hermes is None:
|
|
raise OrchestrationError("Hermes runtime is not configured")
|
|
job = self.hermes.get(result.job_id, tenant_id=result.tenant_id)
|
|
if result.state in {HermesJobState.COMPLETED, HermesJobState.FAILED}:
|
|
if not result.worker_id:
|
|
raise OrchestrationError("terminal Hermes result requires worker_id")
|
|
if job.assigned_worker_id != result.worker_id:
|
|
raise OrchestrationError("Hermes result worker does not own job lease")
|
|
if (
|
|
dispatch.expected_worker_id is not None
|
|
and dispatch.expected_worker_id != result.worker_id
|
|
):
|
|
raise OrchestrationError("Hermes result worker mismatch")
|
|
result_digest = digest_payload(
|
|
{
|
|
"job_id": result.job_id,
|
|
"worker_id": result.worker_id,
|
|
"state": result.state.value,
|
|
"output": result.output,
|
|
"error": result.error,
|
|
"finished_at": result.finished_at,
|
|
}
|
|
)
|
|
if dispatch.state in TERMINAL_DISPATCH_STATES:
|
|
if dispatch.result_digest != result_digest:
|
|
raise DispatchConflict("terminal dispatch received a different result")
|
|
return self.control_plane.get_task(context, dispatch.task_id)
|
|
|
|
task = self.control_plane.get_task(context, dispatch.task_id)
|
|
if task.state != TaskState.RUNNING:
|
|
raise OrchestrationError("task is not RUNNING for active dispatch")
|
|
if result.state == HermesJobState.COMPLETED:
|
|
target = TaskState.COMPLETED
|
|
values = {"result": result.output or {}}
|
|
dispatch_state = DispatchState.COMPLETED
|
|
elif result.state == HermesJobState.FAILED:
|
|
target = TaskState.FAILED
|
|
values = {"error": result.error or {"code": "hermes_failed"}}
|
|
dispatch_state = DispatchState.FAILED
|
|
elif result.state == HermesJobState.CANCELLED:
|
|
target = TaskState.CANCELLED
|
|
values = {}
|
|
dispatch_state = DispatchState.CANCELLED
|
|
else:
|
|
raise OrchestrationError("Hermes reconcile only accepts terminal results")
|
|
|
|
updated_task = self.control_plane.transition_task(
|
|
context,
|
|
task_id=task.task_id,
|
|
target_state=target,
|
|
expected_revision=task.revision,
|
|
**values,
|
|
)
|
|
self.dispatch_repository.update(
|
|
dispatch.dispatch_id,
|
|
tenant_id=dispatch.tenant_id,
|
|
state=dispatch_state,
|
|
expected_worker_id=result.worker_id,
|
|
result_digest=result_digest,
|
|
)
|
|
return updated_task
|
|
|
|
def dispatch_snapshot(self, *, tenant_id: str) -> Mapping[str, DispatchRecord]:
|
|
return {
|
|
item.task_id: item
|
|
for item in self.dispatch_repository.list(tenant_id=tenant_id)
|
|
}
|