clm-kernel 0.0.0__py3-none-any.whl

clm/layer1/pocketflow_scheduler.py

@@ -0,0 +1,519 @@
+"""Pocket Flow Graph Scheduling & Monadic Lifecycle (DEL-388-01, INV-388-01..04).
+
+Implements the category-theoretic Pocket Flow execution lifecycle:
+    Phase 1: prep() — Applicative Functor (<*>), parallel precondition gates
+    Phase 2: exec() — Polynomial Functor (A -> B), isolated computation
+    Phase 3: post() — Monadic Bind (>>=), output-dependent invariant checks with rollback
+
+Enforces Noetherian token conservation (x^T · D = 0) on every transition firing
+and raises NoetherErr on conservation violation.
+"""
+from __future__ import annotations
+
+import asyncio
+import copy
+import logging
+from dataclasses import asdict, dataclass, field
+from typing import Any, Callable
+
+from clm.layer0.exceptions import CLMError
+from clm.layer1.petri import NoetherErr
+from clm.layer1.store import PTRSharedStore
+
+logger = logging.getLogger(__name__)
+
+
+class PostGateRollbackError(CLMError):
+    """Raised when Monadic post-gate verification collapses to ⊥, triggering rollback."""
+    pass
+
+
+@dataclass
+class PrepResult:
+    """Result of the Applicative Pre-Gate (<*>)."""
+    success: bool
+    error: str | None = None
+    context: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class ExecResult:
+    """Result of the Polynomial Functor execution (A -> B)."""
+    success: bool
+    output: Any = None
+    error: str | None = None
+
+
+@dataclass
+class PostResult:
+    """Result of the Monadic Post-Gate (>>=)."""
+    success: bool
+    error: str | None = None
+    witness: Any = None
+
+
+class PrepPhase:
+    """Applicative Pre-Gate (<*>).
+
+    Evaluates independent preconditions concurrently without observing
+    or depending on unproduced outputs.
+    """
+
+    def __init__(
+        self,
+        checks: list[Callable[[Any, dict[str, Any]], bool | tuple[bool, str]]]
+        | Callable[[Any, dict[str, Any]], bool | tuple[bool, str]]
+        | None = None,
+    ):
+        if checks is None:
+            self.checks = []
+        elif isinstance(checks, list):
+            self.checks = checks
+        else:
+            self.checks = [checks]
+
+    def add_check(self, check: Callable[[Any, dict[str, Any]], bool | tuple[bool, str]]) -> None:
+        self.checks.append(check)
+
+    def evaluate(self, store: Any, context: dict[str, Any] | None = None) -> PrepResult:
+        ctx = context or {}
+        for check in self.checks:
+            try:
+                res = check(store, ctx)
+                if isinstance(res, tuple):
+                    passed, msg = res
+                    if not passed:
+                        return PrepResult(success=False, error=msg, context=ctx)
+                elif not res:
+                    return PrepResult(success=False, error="Precondition check failed", context=ctx)
+            except Exception as e:
+                return PrepResult(success=False, error=f"Precondition exception: {e}", context=ctx)
+        return PrepResult(success=True, context=ctx)
+
+    async def evaluate_async(self, store: Any, context: dict[str, Any] | None = None) -> PrepResult:
+        """Evaluate all preconditions concurrently as pure Applicative Functor (<*>)."""
+        ctx = context or {}
+        if not self.checks:
+            return PrepResult(success=True, context=ctx)
+
+        loop = asyncio.get_running_loop()
+
+        def _run_sync(c):
+            return c(store, ctx)
+
+        tasks = [loop.run_in_executor(None, _run_sync, c) for c in self.checks]
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        for res in results:
+            if isinstance(res, Exception):
+                return PrepResult(success=False, error=str(res), context=ctx)
+            if isinstance(res, tuple):
+                passed, msg = res
+                if not passed:
+                    return PrepResult(success=False, error=msg, context=ctx)
+            elif not res:
+                return PrepResult(success=False, error="Precondition check failed", context=ctx)
+
+        return PrepResult(success=True, context=ctx)
+
+
+class ExecPhase:
+    """Polynomial Functor (A -> B).
+
+    Morphism computing isolated output from inputs within the declared context.
+    """
+
+    def __init__(self, execute_fn: Callable[[Any, Any, dict[str, Any]], Any] | None = None):
+        self.execute_fn = execute_fn
+
+    def evaluate(self, inputs: Any, store: Any, context: dict[str, Any] | None = None) -> ExecResult:
+        ctx = context or {}
+        if self.execute_fn is None:
+            return ExecResult(success=True, output=inputs)
+        try:
+            output = self.execute_fn(inputs, store, ctx)
+            return ExecResult(success=True, output=output)
+        except Exception as e:
+            logger.error(f"[ExecPhase] Computation error: {e}", exc_info=True)
+            return ExecResult(success=False, error=str(e))
+
+
+class PostPhase:
+    """Monadic Post-Gate (>>=).
+
+    Output-dependent invariant inspection. If verification fails (Meet = ⊥),
+    triggers transactional rollback.
+    """
+
+    def __init__(self, verify_fn: Callable[[Any, Any, dict[str, Any]], bool | tuple[bool, str]] | None = None):
+        self.verify_fn = verify_fn
+
+    def evaluate(self, output: Any, store: Any, context: dict[str, Any] | None = None) -> PostResult:
+        ctx = context or {}
+        if self.verify_fn is None:
+            return PostResult(success=True, witness={"status": "default_pass"})
+        try:
+            res = self.verify_fn(output, store, ctx)
+            if isinstance(res, tuple):
+                passed, msg = res
+                if not passed:
+                    return PostResult(success=False, error=msg)
+                return PostResult(success=True, witness={"message": msg})
+            elif not res:
+                return PostResult(success=False, error="Monadic post-condition check collapsed to ⊥")
+            return PostResult(success=True, witness={"verified": True})
+        except Exception as e:
+            return PostResult(success=False, error=f"Postcondition exception: {e}")
+
+
+@dataclass
+class DAGNode:
+    """A node in the Pocket Flow DAG, isomorphic to a CPN Transition (t in T)."""
+    id: str
+    prep: PrepPhase | Callable | list[Callable] | None = None
+    exec: ExecPhase | Callable | None = None
+    post: PostPhase | Callable | None = None
+    dependencies: list[str] = field(default_factory=list)
+    inputs: list[str] = field(default_factory=list)
+    outputs: list[str] = field(default_factory=list)
+    tokens_consumed: int = 1
+    tokens_produced: int = 1
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self):
+        # Normalize prep
+        if not isinstance(self.prep, PrepPhase):
+            self.prep = PrepPhase(self.prep)
+        # Normalize exec
+        if not isinstance(self.exec, ExecPhase):
+            self.exec = ExecPhase(self.exec)
+        # Normalize post
+        if not isinstance(self.post, PostPhase):
+            self.post = PostPhase(self.post)
+
+
+@dataclass
+class SchedulerResult:
+    """Aggregated result from executing a Pocket Flow DAG."""
+    success: bool
+    node_results: dict[str, Any] = field(default_factory=dict)
+    store: Any = None
+    error: str | None = None
+    telemetry: dict[str, Any] = field(default_factory=dict)
+
+
+class PocketFlowScheduler:
+    """Coordinates multi-node agentic DAG execution using the Pocket Flow lifecycle.
+
+    Guarantees:
+    - INV-388-01: Applicative prep (<*>) and Monadic post (>>=) typeclass asymmetry.
+    - INV-388-02: Noetherian token conservation (x^T · D = 0^T) across transition firings.
+    - Transactional rollback on post-phase failure.
+    """
+
+    def __init__(self, nodes: list[DAGNode] | None = None):
+        self.nodes: dict[str, DAGNode] = {}
+        if nodes:
+            for n in nodes:
+                self.add_node(n)
+
+    def add_node(self, node: DAGNode) -> None:
+        self.nodes[node.id] = node
+
+    def add_nodes(self, nodes: list[DAGNode]) -> None:
+        for n in nodes:
+            self.add_node(n)
+
+    def topological_sort(self) -> list[DAGNode]:
+        """Compute topological order of DAG nodes."""
+        in_degree: dict[str, int] = dict.fromkeys(self.nodes, 0)
+        adj: dict[str, list[str]] = {nid: [] for nid in self.nodes}
+
+        for nid, node in self.nodes.items():
+            for dep in node.dependencies:
+                if dep in self.nodes:
+                    adj[dep].append(nid)
+                    in_degree[nid] += 1
+
+        queue = [nid for nid, deg in in_degree.items() if deg == 0]
+        order = []
+
+        while queue:
+            curr = queue.pop(0)
+            order.append(self.nodes[curr])
+            for neighbor in adj[curr]:
+                in_degree[neighbor] -= 1
+                if in_degree[neighbor] == 0:
+                    queue.append(neighbor)
+
+        if len(order) != len(self.nodes):
+            raise CLMError(f"PocketFlowScheduler cycle detected: graph contains a cycle among {[nid for nid, deg in in_degree.items() if deg > 0]}")
+
+        return order
+
+    def _snapshot_store(self, store: Any) -> Any:
+        """Create an isolated snapshot of store state for atomic rollback."""
+        if isinstance(store, PTRSharedStore):
+            snap = PTRSharedStore()
+            snap.inputs = copy.deepcopy(store.inputs)
+            snap.outputs = copy.deepcopy(store.outputs)
+            snap.intermediate = copy.deepcopy(store.intermediate)
+            snap.auth_tokens = copy.deepcopy(store.auth_tokens)
+            snap.caller_did = store.caller_did
+            snap.target_did = store.target_did
+            snap.gas_budget = store.gas_budget
+            snap.gas_used = store.gas_used
+            snap.pcard_hash = store.pcard_hash
+            snap.target_hash = store.target_hash
+            return snap
+        elif isinstance(store, dict):
+            return copy.deepcopy(store)
+        return copy.deepcopy(store)
+
+    def _restore_store(self, target_store: Any, snapshot: Any) -> None:
+        """Restore store from snapshot."""
+        if isinstance(target_store, PTRSharedStore) and isinstance(snapshot, PTRSharedStore):
+            target_store.inputs = copy.deepcopy(snapshot.inputs)
+            target_store.outputs = copy.deepcopy(snapshot.outputs)
+            target_store.intermediate = copy.deepcopy(snapshot.intermediate)
+            target_store.auth_tokens = copy.deepcopy(snapshot.auth_tokens)
+            target_store.caller_did = snapshot.caller_did
+            target_store.target_did = snapshot.target_did
+            target_store.gas_budget = snapshot.gas_budget
+            target_store.gas_used = snapshot.gas_used
+        elif isinstance(target_store, dict) and isinstance(snapshot, dict):
+            target_store.clear()
+            target_store.update(copy.deepcopy(snapshot))
+
+    def _get_store_val(self, store: Any, key: str) -> Any:
+        if isinstance(store, PTRSharedStore) or isinstance(store, dict):
+            return store.get(key)
+        return None
+
+    def _set_store_val(self, store: Any, key: str, val: Any) -> None:
+        if isinstance(store, PTRSharedStore):
+            store.set(key, val)
+        elif isinstance(store, dict):
+            store[key] = val
+
+    def verify_noetherian_invariants(
+        self, node: DAGNode, token_markings: dict[str, int]
+    ) -> None:
+        """Assert Noetherian token conservation (x^T · D = 0) for transition firing.
+
+        Raises NoetherErr if input tokens are insufficient or consumed != produced.
+        """
+        # 1. Input token availability check: M >= D^-
+        for inp in node.inputs:
+            current = token_markings.get(inp, 0)
+            if current < node.tokens_consumed:
+                raise NoetherErr(
+                    f"Noetherian Invariant Violation on transition '{node.id}': "
+                    f"Place '{inp}' has {current} tokens, but transition requires {node.tokens_consumed}."
+                )
+
+        # 2. Token conservation balance check: sum(D) = 0 for standard closed transitions
+        delta = node.tokens_produced - node.tokens_consumed
+        if delta != 0 and node.metadata.get("enforce_conservation", True):
+            raise NoetherErr(
+                f"Noetherian Invariant Violation on transition '{node.id}': "
+                f"Conservation x^T · D = 0 violated (consumed {node.tokens_consumed}, produced {node.tokens_produced}, delta = {delta})."
+            )
+
+    def execute(
+        self, store: Any = None, context: dict[str, Any] | None = None
+    ) -> SchedulerResult:
+        """Synchronously execute the Pocket Flow DAG."""
+        if store is None:
+            store = PTRSharedStore()
+        ctx = context or {}
+        order = self.topological_sort()
+
+        # Initialize Petri Net token markings from store inputs
+        token_markings: dict[str, int] = {}
+        for nid, node in self.nodes.items():
+            for inp in node.inputs:
+                if inp not in token_markings:
+                    val = self._get_store_val(store, f"{inp}_tokens")
+                    if val is None:
+                        val = self._get_store_val(store, f"{inp}")
+                    if val is None:
+                        val = 1
+                    token_markings[inp] = int(val) if isinstance(val, (int, float, str)) and str(val).isdigit() else 1
+            for outp in node.outputs:
+                if outp not in token_markings:
+                    val = self._get_store_val(store, f"{outp}_tokens")
+                    if val is None:
+                        val = 0
+                    token_markings[outp] = int(val) if isinstance(val, (int, float, str)) and str(val).isdigit() else 0
+
+        node_results: dict[str, Any] = {}
+
+        for node in order:
+            snapshot = self._snapshot_store(store)
+            logger.debug(f"[PocketFlowScheduler] Firing node '{node.id}'")
+
+            # Invariant check: Noetherian token conservation prior to execution
+            self.verify_noetherian_invariants(node, token_markings)
+
+            # Phase 1: prep (Applicative Functor <*>)
+            prep_res = node.prep.evaluate(store, ctx)
+            if not prep_res.success:
+                err_msg = f"Node '{node.id}' prep phase failed: {prep_res.error}"
+                logger.warning(err_msg)
+                return SchedulerResult(
+                    success=False,
+                    node_results=node_results,
+                    store=store,
+                    error=err_msg,
+                )
+
+            # Resolve input token payload from store inputs
+            input_payload = None
+            if node.inputs:
+                input_payload = {inp: self._get_store_val(store, inp) for inp in node.inputs}
+                if len(node.inputs) == 1:
+                    input_payload = input_payload[node.inputs[0]]
+
+            # Phase 2: exec (Polynomial Functor A -> B)
+            exec_res = node.exec.evaluate(input_payload, store, ctx)
+            if not exec_res.success:
+                self._restore_store(store, snapshot)
+                err_msg = f"Node '{node.id}' exec phase failed: {exec_res.error}"
+                logger.warning(err_msg)
+                return SchedulerResult(
+                    success=False,
+                    node_results=node_results,
+                    store=store,
+                    error=err_msg,
+                )
+
+            # Phase 3: post (Monadic Bind >>=)
+            post_res = node.post.evaluate(exec_res.output, store, ctx)
+            if not post_res.success:
+                # Atomically rollback store to pre-node state (<1ms)
+                self._restore_store(store, snapshot)
+                err_msg = f"Node '{node.id}' post-gate failed (Meet = ⊥): {post_res.error}"
+                logger.warning(f"[ROLLBACK] {err_msg}")
+                return SchedulerResult(
+                    success=False,
+                    node_results=node_results,
+                    store=store,
+                    error=err_msg,
+                )
+
+            # Update token markings upon successful firing: M' = M + D · y
+            for inp in node.inputs:
+                token_markings[inp] -= node.tokens_consumed
+            for outp in node.outputs:
+                token_markings[outp] += node.tokens_produced
+                # Commit output value to store
+                self._set_store_val(store, outp, exec_res.output)
+
+            # Record node execution result
+            node_results[node.id] = {
+                "prep": asdict(prep_res),
+                "exec": asdict(exec_res),
+                "post": asdict(post_res),
+                "output": exec_res.output,
+            }
+
+        return SchedulerResult(
+            success=True,
+            node_results=node_results,
+            store=store,
+            telemetry={"executed_nodes": len(order), "final_markings": token_markings},
+        )
+
+    async def execute_async(
+        self, store: Any = None, context: dict[str, Any] | None = None
+    ) -> SchedulerResult:
+        """Asynchronously execute the Pocket Flow DAG with parallel Applicative prep."""
+        if store is None:
+            store = PTRSharedStore()
+        ctx = context or {}
+        order = self.topological_sort()
+
+        token_markings: dict[str, int] = {}
+        for nid, node in self.nodes.items():
+            for inp in node.inputs:
+                if inp not in token_markings:
+                    val = self._get_store_val(store, f"{inp}_tokens")
+                    if val is None:
+                        val = self._get_store_val(store, f"{inp}")
+                    if val is None:
+                        val = 1
+                    token_markings[inp] = int(val) if isinstance(val, (int, float, str)) and str(val).isdigit() else 1
+            for outp in node.outputs:
+                if outp not in token_markings:
+                    val = self._get_store_val(store, f"{outp}_tokens")
+                    if val is None:
+                        val = 0
+                    token_markings[outp] = int(val) if isinstance(val, (int, float, str)) and str(val).isdigit() else 0
+
+        node_results: dict[str, Any] = {}
+
+        for node in order:
+            snapshot = self._snapshot_store(store)
+            self.verify_noetherian_invariants(node, token_markings)
+
+            # Phase 1: Async Applicative prep (<*>)
+            prep_res = await node.prep.evaluate_async(store, ctx)
+            if not prep_res.success:
+                err_msg = f"Node '{node.id}' prep phase failed: {prep_res.error}"
+                return SchedulerResult(success=False, node_results=node_results, store=store, error=err_msg)
+
+            input_payload = None
+            if node.inputs:
+                input_payload = {inp: self._get_store_val(store, inp) for inp in node.inputs}
+                if len(node.inputs) == 1:
+                    input_payload = input_payload[node.inputs[0]]
+
+            # Phase 2: exec
+            loop = asyncio.get_running_loop()
+            exec_res = await loop.run_in_executor(None, node.exec.evaluate, input_payload, store, ctx)
+            if not exec_res.success:
+                self._restore_store(store, snapshot)
+                err_msg = f"Node '{node.id}' exec phase failed: {exec_res.error}"
+                return SchedulerResult(success=False, node_results=node_results, store=store, error=err_msg)
+
+            # Phase 3: post
+            post_res = node.post.evaluate(exec_res.output, store, ctx)
+            if not post_res.success:
+                self._restore_store(store, snapshot)
+                err_msg = f"Node '{node.id}' post-gate failed (Meet = ⊥): {post_res.error}"
+                return SchedulerResult(success=False, node_results=node_results, store=store, error=err_msg)
+
+            for inp in node.inputs:
+                token_markings[inp] -= node.tokens_consumed
+            for outp in node.outputs:
+                token_markings[outp] += node.tokens_produced
+                self._set_store_val(store, outp, exec_res.output)
+
+            node_results[node.id] = {
+                "prep": asdict(prep_res),
+                "exec": asdict(exec_res),
+                "post": asdict(post_res),
+                "output": exec_res.output,
+            }
+
+        return SchedulerResult(
+            success=True,
+            node_results=node_results,
+            store=store,
+            telemetry={"executed_nodes": len(order), "final_markings": token_markings},
+        )
+
+
+__all__ = [
+    "PrepResult",
+    "ExecResult",
+    "PostResult",
+    "PrepPhase",
+    "ExecPhase",
+    "PostPhase",
+    "DAGNode",
+    "PocketFlowScheduler",
+    "SchedulerResult",
+    "PostGateRollbackError",
+]