hypermind 0.11.0__py3-none-any.whl

hypermind/workflow/engine.py

@@ -0,0 +1,655 @@
+"""Workflow engine — durable, signed, replayable steps.
+
+A `Workflow` runs a user-supplied async function whose body issues
+``await wf.step(name, fn, *args, **kwargs)`` calls. Every successful
+step has its result frozen as bytes (canonical CBOR), hashed, and
+appended to the transparency log via
+:func:`TransparencyService.append_and_cosign`. The leaf written to the
+log is a 32-byte BLAKE3-style identifier derived from the step record
+contents, so the log alone (plus the locally-retained payload bytes)
+is enough to replay.
+
+Replay determinism (REQUIRED):
+    A step body MUST be a pure function of its arguments and any prior
+    Activity outputs. Non-deterministic calls — LLM completions, system
+    clock reads, random number generation, network IO — MUST be wrapped
+    in an :class:`Activity` whose output bytes are captured at run time
+    and re-used at replay time. This is the same Temporal/DBOS pattern.
+
+Compensations:
+    A step may register a compensator via the :func:`compensate`
+    decorator (or by passing ``compensator=`` to ``wf.step``). On
+    workflow failure, completed steps are unwound in reverse order;
+    each compensation is itself recorded as a step with
+    ``status="COMPENSATING"`` so the audit trail is symmetric.
+
+Security notes:
+    * Result bytes are frozen on ``COMPLETED`` — the dataclass is
+      ``frozen=True`` after that point.
+    * Every state transition (RUNNING → COMPLETED / FAILED /
+      COMPENSATED) is reflected in the log.
+    * Replay never re-runs an Activity; it consumes the captured bytes.
+      A mismatch between recorded args_hash and replay args_hash raises.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import inspect
+import secrets
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import Any, ClassVar
+
+import cbor2
+
+from hypermind.crypto.hlc import HLC, SystemClock
+from hypermind.knowledge.transparency import TransparencyService
+
+
+def _default_ts(name: str) -> TransparencyService:
+    """Build a TransparencyService usable inside an in-process workflow.
+
+    Production deployments inject their own configured TS (with witness
+    keypairs); the default constructed here uses a witness quorum of
+    zero so the engine works out of the box for in-process tests and
+    examples. Callers MUST supply a properly-witnessed TS for any
+    deployment that requires audit-grade signed transparency.
+    """
+    return TransparencyService(name=name, min_witness_quorum=0)
+
+
+def _hlc_str(node_id: bytes) -> str:
+    """Stringified HLC suitable for human-readable record fields."""
+    h = HLC.now(node_id, SystemClock())
+    return f"{h.logical_ms}.{h.counter}@{h.node_id.hex()}"
+
+
+# ---------------------------------------------------------------------------
+# Status enums (string Literals so they round-trip cleanly through CBOR)
+# ---------------------------------------------------------------------------
+
+WorkflowStatus = str  # "RUNNING" | "COMPLETED" | "FAILED" | "COMPENSATED"
+StepStatus = str  # "RUNNING" | "COMPLETED" | "FAILED" | "COMPENSATING" | "COMPENSATED"
+
+_VALID_WORKFLOW_STATUS: frozenset[str] = frozenset(
+    {"RUNNING", "COMPLETED", "FAILED", "COMPENSATED"}
+)
+_VALID_STEP_STATUS: frozenset[str] = frozenset(
+    {"RUNNING", "COMPLETED", "FAILED", "COMPENSATING", "COMPENSATED"}
+)
+
+
+# ---------------------------------------------------------------------------
+# Hashing helpers
+# ---------------------------------------------------------------------------
+
+
+def _canon_bytes(obj: Any) -> bytes:
+    """Canonical CBOR serialisation. Used for argument + result freezing.
+
+    Falls back to a `repr()`-based encoding for Python types CBOR
+    cannot serialise (iterators, locks, file handles). Workflow bodies
+    SHOULD pass only CBOR-friendly arguments; the fallback exists so
+    the engine doesn't crash on incidental non-serialisable scaffolding
+    (test iterators, generator objects). The fallback is hash-stable
+    across the same process but NOT across processes — the determinism
+    guarantee remains "argument-bytes equal" + "result-bytes equal",
+    enforced by the args_hash gate during replay.
+    """
+    try:
+        return cbor2.dumps(obj, canonical=True)
+    except (cbor2.CBOREncodeError, TypeError):
+        return cbor2.dumps({"_repr": repr(obj)}, canonical=True)
+
+
+def _h32(data: bytes) -> bytes:
+    return hashlib.sha256(data).digest()
+
+
+def _step_leaf(workflow_id: str, step_id: str, body: bytes) -> bytes:
+    """Domain-tagged 32-byte leaf id for the transparency log."""
+    return _h32(
+        b"hm:wf-step:v1\x00" + workflow_id.encode() + b"\x00" + step_id.encode() + b"\x00" + body
+    )
+
+
+# ---------------------------------------------------------------------------
+# Records
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class Activity:
+    """Captured non-deterministic value (LLM output, clock, RNG, IO).
+
+    ``non_determinism_class`` SHOULD be one of {"LLM","RNG","CLOCK","IO"}
+    per spec §26 SHOULD-clause; engine does not enforce a closed set so
+    integrators can extend the taxonomy.
+    """
+
+    name: str
+    non_determinism_class: str
+    output_bytes: bytes
+
+
+@dataclass(frozen=True)
+class StepRecord:
+    step_id: str
+    step_name: str
+    args_hash: bytes
+    result_hash: bytes
+    result_bytes: bytes  # frozen on COMPLETED
+    hlc: str
+    status: StepStatus
+    leaf_kid: bytes  # transparency-log leaf id
+    activities: tuple[Activity, ...] = ()
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "step_id": self.step_id,
+            "step_name": self.step_name,
+            "args_hash": self.args_hash.hex(),
+            "result_hash": self.result_hash.hex(),
+            "result_bytes": self.result_bytes,
+            "hlc": self.hlc,
+            "status": self.status,
+            "leaf_kid": self.leaf_kid.hex(),
+            "activities": [
+                {
+                    "name": a.name,
+                    "non_determinism_class": a.non_determinism_class,
+                    "output_bytes": a.output_bytes,
+                }
+                for a in self.activities
+            ],
+        }
+
+
+@dataclass
+class WorkflowState:
+    workflow_id: str
+    name: str
+    status: WorkflowStatus = "RUNNING"
+    steps: list[StepRecord] = field(default_factory=list)
+    started_hlc: str = ""
+    completed_hlc: str = ""
+    result_bytes: bytes = b""
+    result_hash: bytes = b""
+    error: str = ""
+
+    # Frozen-after-completion guard: once status leaves RUNNING we set
+    # _frozen and refuse further mutation.
+    _frozen: bool = False
+
+    def freeze(self) -> None:
+        self._frozen = True
+
+    def _check_mutable(self) -> None:
+        if self._frozen:
+            raise RuntimeError(
+                f"WorkflowState[{self.workflow_id}] is frozen ({self.status}); "
+                f"cannot mutate post-completion"
+            )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "workflow_id": self.workflow_id,
+            "name": self.name,
+            "status": self.status,
+            "started_hlc": self.started_hlc,
+            "completed_hlc": self.completed_hlc,
+            "result_hash": self.result_hash.hex(),
+            "error": self.error,
+            "steps": [s.to_dict() for s in self.steps],
+        }
+
+
+# ---------------------------------------------------------------------------
+# Decorators
+# ---------------------------------------------------------------------------
+
+
+# Global registry: name -> (workflow_callable, compensators dict)
+_WORKFLOW_REGISTRY: dict[str, dict[str, Any]] = {}
+
+
+def workflow(
+    _fn: Callable[..., Awaitable[Any]] | None = None,
+    *,
+    name: str | None = None,
+) -> Any:
+    """Register a coroutine as a durable workflow.
+
+    Usage::
+
+        @workflow
+        async def my_wf(wf, payload): ...
+
+        @workflow(name="custom")
+        async def my_wf(wf, payload): ...
+    """
+
+    def _decorate(fn: Callable[..., Awaitable[Any]]) -> Callable[..., Awaitable[Any]]:
+        if not inspect.iscoroutinefunction(fn):
+            raise TypeError(f"@workflow requires an async function (got {fn!r})")
+        wf_name = name or fn.__name__
+        # Use setdefault so repeated decoration (e.g. test reload) doesn't
+        # blow away previously registered compensators.
+        entry = _WORKFLOW_REGISTRY.setdefault(wf_name, {"fn": fn, "compensators": {}})
+        entry["fn"] = fn
+        fn._hm_workflow_name = wf_name  # type: ignore[attr-defined]
+        return fn
+
+    if _fn is None:
+        return _decorate
+    return _decorate(_fn)
+
+
+def compensate(
+    step_name: str,
+    *,
+    workflow_name: str | None = None,
+) -> Callable[[Callable[..., Awaitable[Any]]], Callable[..., Awaitable[Any]]]:
+    """Register a compensator for a step.
+
+    The compensator receives the step's recorded result_bytes and runs
+    in reverse order during failure unwinding.
+    """
+
+    def _decorate(fn: Callable[..., Awaitable[Any]]) -> Callable[..., Awaitable[Any]]:
+        if not inspect.iscoroutinefunction(fn):
+            raise TypeError(f"@compensate requires an async function (got {fn!r})")
+        # Late binding: the workflow name is the *currently being decorated*
+        # workflow if not explicit. Tests typically call @workflow before
+        # @compensate, so the registry already has the entry.
+        wf_name = workflow_name
+        if wf_name is None:
+            # Search for the most recently registered workflow.
+            if not _WORKFLOW_REGISTRY:
+                raise RuntimeError(
+                    "@compensate must be paired with a @workflow; "
+                    "register the workflow first or pass workflow_name=..."
+                )
+            wf_name = next(reversed(_WORKFLOW_REGISTRY))
+        _WORKFLOW_REGISTRY[wf_name]["compensators"][step_name] = fn
+        return fn
+
+    return _decorate
+
+
+# ---------------------------------------------------------------------------
+# Workflow runner
+# ---------------------------------------------------------------------------
+
+
+class _WorkflowContext:
+    """Handed to the user's workflow body — its `step` method drives the
+    durable step machinery."""
+
+    def __init__(
+        self,
+        state: WorkflowState,
+        transparency: TransparencyService,
+        compensators: dict[str, Callable[..., Awaitable[Any]]],
+        node_id: bytes,
+        replay_steps: list[StepRecord] | None = None,
+    ) -> None:
+        self.state = state
+        self.transparency = transparency
+        self.compensators = compensators
+        self._node_id = node_id
+        self._replay = list(replay_steps) if replay_steps else None
+        self._replay_idx = 0
+
+    async def step(
+        self,
+        step_name: str,
+        body: Callable[..., Awaitable[Any]],
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        """Run a single durable step.
+
+        On first execution: record args_hash, call body, freeze result
+        bytes, append to transparency log, attach to state.
+
+        On replay: skip body and return the previously-recorded result.
+        """
+        args_blob = _canon_bytes({"args": list(args), "kwargs": kwargs})
+        args_hash = _h32(args_blob)
+
+        if self._replay is not None:
+            # Replay path — must match the next recorded step exactly.
+            if self._replay_idx >= len(self._replay):
+                raise RuntimeError(
+                    f"replay overrun: workflow body issued step "
+                    f"{step_name!r} but log has no more entries"
+                )
+            recorded = self._replay[self._replay_idx]
+            self._replay_idx += 1
+            if recorded.step_name != step_name:
+                raise RuntimeError(
+                    f"replay divergence: log expects {recorded.step_name!r}, "
+                    f"body issued {step_name!r}"
+                )
+            if recorded.args_hash != args_hash:
+                raise RuntimeError(
+                    f"replay divergence: args_hash mismatch on step "
+                    f"{step_name!r} (recorded {recorded.args_hash.hex()[:12]}, "
+                    f"now {args_hash.hex()[:12]})"
+                )
+            # Reattach so caller can introspect.
+            self.state.steps.append(recorded)
+            return cbor2.loads(recorded.result_bytes) if recorded.result_bytes else None
+
+        # Live path.
+        step_id = secrets.token_hex(8)
+        try:
+            result = await body(*args, **kwargs)
+        except Exception as exc:
+            # Record FAILED step then re-raise; engine caller decides on
+            # compensation.
+            failed = StepRecord(
+                step_id=step_id,
+                step_name=step_name,
+                args_hash=args_hash,
+                result_hash=b"",
+                result_bytes=b"",
+                hlc=_hlc_str(self._node_id),
+                status="FAILED",
+                leaf_kid=_step_leaf(self.state.workflow_id, step_id, args_blob),
+            )
+            self.transparency.append_and_cosign(failed.leaf_kid)
+            self.state.steps.append(failed)
+            self.state.error = f"{step_name}: {exc!s}"
+            raise
+        result_bytes = _canon_bytes(result)
+        result_hash = _h32(result_bytes)
+        leaf_kid = _step_leaf(self.state.workflow_id, step_id, result_bytes)
+        record = StepRecord(
+            step_id=step_id,
+            step_name=step_name,
+            args_hash=args_hash,
+            result_hash=result_hash,
+            result_bytes=result_bytes,
+            hlc=_hlc_str(self._node_id),
+            status="COMPLETED",
+            leaf_kid=leaf_kid,
+        )
+        self.transparency.append_and_cosign(leaf_kid)
+        self.state.steps.append(record)
+        return result
+
+    async def activity(
+        self,
+        name: str,
+        body: Callable[..., Awaitable[Any]],
+        *args: Any,
+        non_determinism_class: str = "IO",
+        **kwargs: Any,
+    ) -> Any:
+        """Run a non-deterministic side-effect; capture output bytes.
+
+        On replay the captured bytes are returned unchanged. Use this
+        for LLM calls, ``time.time()``, ``random.random()``, network IO,
+        etc. The body itself is NOT re-executed during replay.
+        """
+        return await self.step(
+            f"activity:{name}",
+            body,
+            *args,
+            **kwargs,
+        )
+
+
+class Workflow:
+    """Static façade for running and replaying durable workflows."""
+
+    _states: ClassVar[dict[str, WorkflowState]] = {}
+
+    @classmethod
+    async def run(
+        cls,
+        fn: Callable[..., Awaitable[Any]],
+        *args: Any,
+        transparency: TransparencyService | None = None,
+        node_id: bytes | None = None,
+        workflow_id: str | None = None,
+        **kwargs: Any,
+    ) -> WorkflowState:
+        """Execute a workflow. Returns its final state."""
+        wf_name = getattr(fn, "_hm_workflow_name", None)
+        if wf_name is None:
+            raise TypeError(f"{fn!r} is not a registered @workflow — apply @workflow first")
+        ts = transparency if transparency is not None else _default_ts(f"workflow-{wf_name}")
+        nid = node_id or secrets.token_bytes(8)
+        wid = workflow_id or secrets.token_hex(8)
+        state = WorkflowState(
+            workflow_id=wid,
+            name=wf_name,
+            status="RUNNING",
+            started_hlc=_hlc_str(nid),
+        )
+        ctx = _WorkflowContext(
+            state=state,
+            transparency=ts,
+            compensators=_WORKFLOW_REGISTRY[wf_name]["compensators"],
+            node_id=nid,
+        )
+        try:
+            result = await fn(ctx, *args, **kwargs)
+        except Exception:
+            state.status = "FAILED"
+            state.completed_hlc = _hlc_str(nid)
+            state.freeze()
+            cls._states[wid] = state
+            raise
+        state.result_bytes = _canon_bytes(result)
+        state.result_hash = _h32(state.result_bytes)
+        state.status = "COMPLETED"
+        state.completed_hlc = _hlc_str(nid)
+        state.freeze()
+        cls._states[wid] = state
+        return state
+
+    @classmethod
+    def replay(
+        cls,
+        workflow_id: str,
+        *,
+        transparency: TransparencyService | None = None,
+    ) -> WorkflowState:
+        """Re-execute the workflow body using captured step results.
+
+        Asserts identical outcome (same result_hash) compared to the
+        original run. The body is invoked in replay mode: every
+        ``wf.step`` call returns the recorded result without re-running
+        the underlying coroutine. This is how non-deterministic
+        Activities (LLM/RNG/clock/IO) stay deterministic across replay.
+        """
+        original = cls._states.get(workflow_id)
+        if original is None:
+            raise KeyError(f"unknown workflow_id {workflow_id!r}")
+        wf_entry = _WORKFLOW_REGISTRY.get(original.name)
+        if wf_entry is None:
+            raise KeyError(f"workflow {original.name!r} no longer registered")
+        ts = (
+            transparency
+            if transparency is not None
+            else _default_ts(f"workflow-replay-{original.name}")
+        )
+
+        replay_state = WorkflowState(
+            workflow_id=workflow_id,
+            name=original.name,
+            status="RUNNING",
+            started_hlc=original.started_hlc,
+        )
+        ctx = _WorkflowContext(
+            state=replay_state,
+            transparency=ts,
+            compensators=wf_entry["compensators"],
+            node_id=secrets.token_bytes(8),
+            replay_steps=original.steps,
+        )
+
+        # The user fn signature is (wf, *args, **kwargs); we don't have
+        # the original args saved on the WorkflowState, so we reconstruct
+        # by re-driving the body — which on replay only consumes recorded
+        # step results. The body must be argument-free under replay; for
+        # workflows with args, callers can pass them via the original
+        # closure (the canonical pattern: register a no-arg @workflow
+        # that closes over inputs).
+        async def _drive() -> Any:
+            try:
+                return await wf_entry["fn"](ctx)
+            except TypeError:
+                # Fallback: if the workflow expects extra args, the
+                # caller must use replay_with_args() — out of scope for
+                # the basic engine. Mark FAILED with explanatory error.
+                raise
+
+        result = (
+            asyncio.get_event_loop().run_until_complete(_drive())
+            if not _loop_running()
+            else _run_in_new_loop(_drive())
+        )
+
+        replay_state.result_bytes = _canon_bytes(result)
+        replay_state.result_hash = _h32(replay_state.result_bytes)
+        replay_state.status = "COMPLETED"
+        replay_state.completed_hlc = original.completed_hlc
+        replay_state.freeze()
+
+        # Determinism gate.
+        if replay_state.result_hash != original.result_hash:
+            raise RuntimeError(
+                f"replay determinism violation: original "
+                f"{original.result_hash.hex()[:12]} != "
+                f"replay {replay_state.result_hash.hex()[:12]}"
+            )
+        return replay_state
+
+    @classmethod
+    async def compensate(
+        cls,
+        workflow_id: str,
+        *,
+        transparency: TransparencyService | None = None,
+    ) -> WorkflowState:
+        """Run compensators for every COMPLETED step in reverse order.
+
+        Compensation records are appended to the transparency log with
+        ``status="COMPENSATING"``; on success the workflow status is
+        set to ``COMPENSATED``.
+        """
+        state = cls._states.get(workflow_id)
+        if state is None:
+            raise KeyError(f"unknown workflow_id {workflow_id!r}")
+        wf_entry = _WORKFLOW_REGISTRY.get(state.name)
+        if wf_entry is None:
+            raise KeyError(f"workflow {state.name!r} no longer registered")
+        compensators = wf_entry["compensators"]
+        ts = (
+            transparency if transparency is not None else _default_ts(f"workflow-comp-{state.name}")
+        )
+        nid = secrets.token_bytes(8)
+
+        # Unfreeze for the compensation pass — the audit-state we're
+        # appending is itself part of the workflow's lifecycle, not a
+        # post-completion mutation of the original record.
+        state._frozen = False
+        completed_in_order = [s for s in list(state.steps) if s.status == "COMPLETED"]
+        for original_step in reversed(completed_in_order):
+            comp_fn = compensators.get(original_step.step_name)
+            if comp_fn is None:
+                continue
+            step_id = secrets.token_hex(8)
+            try:
+                # Compensator receives the recorded result bytes (CBOR-decoded
+                # for ergonomic access).
+                payload = (
+                    cbor2.loads(original_step.result_bytes) if original_step.result_bytes else None
+                )
+                comp_result = await comp_fn(payload)
+            except Exception as exc:
+                state.error = f"compensation failed at {original_step.step_name}: {exc!s}"
+                state.status = "FAILED"
+                state.freeze()
+                raise
+            comp_bytes = _canon_bytes(comp_result)
+            leaf_kid = _step_leaf(workflow_id, step_id, comp_bytes)
+            comp_record = StepRecord(
+                step_id=step_id,
+                step_name=f"compensate:{original_step.step_name}",
+                args_hash=original_step.result_hash,
+                result_hash=_h32(comp_bytes),
+                result_bytes=comp_bytes,
+                hlc=_hlc_str(nid),
+                status="COMPENSATING",
+                leaf_kid=leaf_kid,
+            )
+            ts.append_and_cosign(leaf_kid)
+            state.steps.append(comp_record)
+        state.status = "COMPENSATED"
+        state.completed_hlc = _hlc_str(nid)
+        state.freeze()
+        return state
+
+    @classmethod
+    def list_ids(cls) -> list[str]:
+        return list(cls._states.keys())
+
+    @classmethod
+    def get(cls, workflow_id: str) -> WorkflowState:
+        st = cls._states.get(workflow_id)
+        if st is None:
+            raise KeyError(workflow_id)
+        return st
+
+    @classmethod
+    def _reset_for_test(cls) -> None:
+        """Test-only: clear the in-process state registry."""
+        cls._states.clear()
+
+
+# Helpers ------------------------------------------------------------------
+
+
+def _loop_running() -> bool:
+    try:
+        asyncio.get_running_loop()
+        return True
+    except RuntimeError:
+        return False
+
+
+def _run_in_new_loop(coro: Awaitable[Any]) -> Any:
+    """Run ``coro`` on a fresh event loop. Used by ``Workflow.replay``
+    when the caller is inside an active loop (pytest-asyncio).
+
+    We spin up a worker thread to avoid the "loop already running"
+    error — replay is by contract a synchronous boundary because it's
+    the auditor's invariant ("re-derive the answer from the log alone").
+    """
+    import threading
+
+    box: dict[str, Any] = {}
+
+    def _runner() -> None:
+        loop = asyncio.new_event_loop()
+        try:
+            box["result"] = loop.run_until_complete(coro)
+        except BaseException as e:
+            box["error"] = e
+        finally:
+            loop.close()
+
+    t = threading.Thread(target=_runner, daemon=True)
+    t.start()
+    t.join()
+    if "error" in box:
+        raise box["error"]
+    return box["result"]