229 lines
8.3 KiB
Python
229 lines
8.3 KiB
Python
"""Signed, tenant-isolated execution adapter for OpenClaw runtimes."""
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
from dataclasses import replace
|
|
from typing import Callable, Mapping, Protocol
|
|
|
|
from .contract import (
|
|
CommandState,
|
|
OpenClawCommand,
|
|
OpenClawResult,
|
|
RuntimeHealth,
|
|
utc_now,
|
|
verify_command_signature,
|
|
)
|
|
|
|
|
|
class AdapterError(RuntimeError):
|
|
"""Base adapter failure."""
|
|
|
|
|
|
class CommandRejected(AdapterError):
|
|
"""A command failed signature or contract validation."""
|
|
|
|
|
|
class CommandNotFound(AdapterError):
|
|
"""A tenant-scoped command is not visible."""
|
|
|
|
|
|
class CommandConflict(AdapterError):
|
|
"""A command or idempotency key conflicts with an existing execution."""
|
|
|
|
|
|
class RuntimeTransport(Protocol):
|
|
def execute(self, command: OpenClawCommand, timeout_seconds: float) -> OpenClawResult: ...
|
|
def cancel(self, command_id: str) -> bool: ...
|
|
def health(self) -> RuntimeHealth: ...
|
|
|
|
|
|
class CallbackSink(Protocol):
|
|
def on_started(self, result: OpenClawResult) -> None: ...
|
|
def on_completed(self, result: OpenClawResult) -> None: ...
|
|
def on_failed(self, result: OpenClawResult) -> None: ...
|
|
|
|
|
|
class InMemoryCallbackSink:
|
|
def __init__(self) -> None:
|
|
self._lock = threading.RLock()
|
|
self.started: list[OpenClawResult] = []
|
|
self.completed: list[OpenClawResult] = []
|
|
self.failed: list[OpenClawResult] = []
|
|
|
|
def on_started(self, result: OpenClawResult) -> None:
|
|
with self._lock:
|
|
self.started.append(result)
|
|
|
|
def on_completed(self, result: OpenClawResult) -> None:
|
|
with self._lock:
|
|
self.completed.append(result)
|
|
|
|
def on_failed(self, result: OpenClawResult) -> None:
|
|
with self._lock:
|
|
self.failed.append(result)
|
|
|
|
|
|
SigningMaterialResolver = Callable[[str, str], bytes | str]
|
|
|
|
|
|
class OpenClawAdapter:
|
|
"""Verify, de-duplicate, execute, cancel and reconcile commands."""
|
|
|
|
def __init__(
|
|
self,
|
|
transport: RuntimeTransport,
|
|
signing_material_resolver: SigningMaterialResolver,
|
|
*,
|
|
callback_sink: CallbackSink | None = None,
|
|
default_timeout_seconds: float = 300.0,
|
|
) -> None:
|
|
if default_timeout_seconds <= 0:
|
|
raise ValueError("default_timeout_seconds must be positive")
|
|
self.transport = transport
|
|
self.signing_material_resolver = signing_material_resolver
|
|
self.callback_sink = callback_sink or InMemoryCallbackSink()
|
|
self.default_timeout_seconds = float(default_timeout_seconds)
|
|
self._lock = threading.RLock()
|
|
self._commands: dict[str, OpenClawCommand] = {}
|
|
self._results: dict[str, OpenClawResult] = {}
|
|
self._idempotency: dict[tuple[str, str], tuple[str, str]] = {}
|
|
|
|
def _verify(self, command: OpenClawCommand) -> None:
|
|
try:
|
|
key_material = self.signing_material_resolver(command.tenant_id, command.key_id)
|
|
except Exception as exc:
|
|
raise CommandRejected("signing material unavailable") from exc
|
|
if not verify_command_signature(command, key_material):
|
|
raise CommandRejected("invalid or expired command signature")
|
|
|
|
@staticmethod
|
|
def _validate_result(command: OpenClawCommand, result: OpenClawResult) -> None:
|
|
if (
|
|
result.command_id != command.command_id
|
|
or result.tenant_id != command.tenant_id
|
|
or result.task_id != command.task_id
|
|
or result.trace_id != command.trace_id
|
|
):
|
|
raise AdapterError("runtime returned mismatched command identity")
|
|
if result.state not in {
|
|
CommandState.COMPLETED,
|
|
CommandState.FAILED,
|
|
CommandState.CANCELLED,
|
|
CommandState.TIMED_OUT,
|
|
}:
|
|
raise AdapterError("runtime returned a non-terminal result")
|
|
|
|
def execute(
|
|
self,
|
|
command: OpenClawCommand,
|
|
*,
|
|
timeout_seconds: float | None = None,
|
|
) -> OpenClawResult:
|
|
timeout = self.default_timeout_seconds if timeout_seconds is None else float(timeout_seconds)
|
|
if timeout <= 0:
|
|
raise ValueError("timeout_seconds must be positive")
|
|
self._verify(command)
|
|
digest = command.signature
|
|
idem_ref = (command.tenant_id, command.idempotency_key)
|
|
|
|
with self._lock:
|
|
prior_command = self._commands.get(command.command_id)
|
|
if prior_command is not None:
|
|
if prior_command.signature != command.signature:
|
|
raise CommandConflict("command_id reused with different content")
|
|
return self._results[command.command_id]
|
|
prior_idem = self._idempotency.get(idem_ref)
|
|
if prior_idem is not None:
|
|
prior_id, prior_digest = prior_idem
|
|
if prior_digest != digest:
|
|
raise CommandConflict("idempotency key reused with different content")
|
|
return self._results[prior_id]
|
|
running = OpenClawResult(
|
|
command_id=command.command_id,
|
|
tenant_id=command.tenant_id,
|
|
task_id=command.task_id,
|
|
trace_id=command.trace_id,
|
|
state=CommandState.RUNNING,
|
|
started_at=utc_now(),
|
|
)
|
|
self._commands[command.command_id] = command
|
|
self._results[command.command_id] = running
|
|
self._idempotency[idem_ref] = (command.command_id, digest)
|
|
|
|
self.callback_sink.on_started(running)
|
|
try:
|
|
result = self.transport.execute(command, timeout)
|
|
self._validate_result(command, result)
|
|
except TimeoutError as exc:
|
|
result = replace(
|
|
running,
|
|
state=CommandState.TIMED_OUT,
|
|
error={"code": "runtime_timeout", "message": str(exc) or "runtime timeout"},
|
|
finished_at=utc_now(),
|
|
)
|
|
except Exception as exc:
|
|
result = replace(
|
|
running,
|
|
state=CommandState.FAILED,
|
|
error={"code": "runtime_error", "message": str(exc)},
|
|
finished_at=utc_now(),
|
|
)
|
|
|
|
if result.started_at is None:
|
|
result = replace(result, started_at=running.started_at)
|
|
if result.finished_at is None:
|
|
result = replace(result, finished_at=utc_now())
|
|
with self._lock:
|
|
self._results[command.command_id] = result
|
|
if result.state == CommandState.COMPLETED:
|
|
self.callback_sink.on_completed(result)
|
|
else:
|
|
self.callback_sink.on_failed(result)
|
|
return result
|
|
|
|
def status(self, command_id: str, *, tenant_id: str) -> OpenClawResult:
|
|
with self._lock:
|
|
command = self._commands.get(command_id)
|
|
result = self._results.get(command_id)
|
|
if command is None or result is None or command.tenant_id != tenant_id:
|
|
raise CommandNotFound("command not found")
|
|
return result
|
|
|
|
def cancel(self, command_id: str, *, tenant_id: str) -> OpenClawResult:
|
|
with self._lock:
|
|
command = self._commands.get(command_id)
|
|
current = self._results.get(command_id)
|
|
if command is None or current is None or command.tenant_id != tenant_id:
|
|
raise CommandNotFound("command not found")
|
|
if current.state in {
|
|
CommandState.COMPLETED,
|
|
CommandState.FAILED,
|
|
CommandState.CANCELLED,
|
|
CommandState.TIMED_OUT,
|
|
CommandState.REJECTED,
|
|
}:
|
|
return current
|
|
if not self.transport.cancel(command_id):
|
|
raise AdapterError("runtime rejected cancellation")
|
|
cancelled = replace(
|
|
current,
|
|
state=CommandState.CANCELLED,
|
|
error={"code": "cancelled", "message": "cancelled by AIOS"},
|
|
finished_at=utc_now(),
|
|
)
|
|
with self._lock:
|
|
self._results[command_id] = cancelled
|
|
self.callback_sink.on_failed(cancelled)
|
|
return cancelled
|
|
|
|
def health(self) -> RuntimeHealth:
|
|
return self.transport.health()
|
|
|
|
def snapshot(self, *, tenant_id: str) -> Mapping[str, OpenClawResult]:
|
|
with self._lock:
|
|
return {
|
|
command_id: result
|
|
for command_id, result in self._results.items()
|
|
if self._commands[command_id].tenant_id == tenant_id
|
|
}
|