feat: CEK Machine — abstract machine interpreter for agent programs

Implements the Agent CEK Machine (Paper10 §5) with:
- 18 transition rules (9 dispatch + 2 yield + 7 return)
- YIELD mechanism for LLM/Tool oracle calls
- Full transition trace with labels and cost vectors
- Support for all 11 lambdagent constructs
- Step-by-step execution and one-shot run modes

10/10 tests passing: Tool, Compose, 3-stage, If, Loop, Pair,
Guard OK/Fail, Compose+Pair, Loop+Compose.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

lambdagent/cek_machine.py

@@ -0,0 +1,639 @@
+"""
+lambdagent.cek_machine — Agent CEK Machine
+
+An abstract machine implementation of the operational semantics
+defined in Paper10. The CEK machine provides:
+
+  1. Step-by-step execution with full state inspection
+  2. YIELD mechanism for LLM/Tool oracle calls
+  3. Complete trace of all transitions (dispatch, yield, return)
+  4. Cost vector tracking at every step
+
+The machine corresponds 1-to-1 with the small-step reduction rules
+in Paper10 §4, and the transition rules in §5.
+
+References:
+  - Felleisen & Friedman (1986): CEK Machine
+  - Paper10 §5: Agent CEK Machine
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from enum import Enum, auto
+
+
+# ============================================================
+# Values and Cost
+# ============================================================
+
+@dataclass(frozen=True)
+class CostVector:
+    """Cost vector c = (tokens, latency_s, money_usd)"""
+    tokens: int = 0
+    latency: float = 0.0
+    money: float = 0.0
+
+    def __add__(self, other: CostVector) -> CostVector:
+        return CostVector(
+            self.tokens + other.tokens,
+            self.latency + other.latency,
+            self.money + other.money,
+        )
+
+    def __repr__(self):
+        return f"Cost(tok={self.tokens}, lat={self.latency:.2f}s, ${self.money:.4f})"
+
+
+ZERO_COST = CostVector()
+
+
+# ============================================================
+# Labels (observable actions)
+# ============================================================
+
+class LabelKind(Enum):
+    TAU = auto()       # silent / internal
+    LLM = auto()       # LLM oracle call
+    TOOL = auto()      # tool oracle call
+    MEM = auto()       # memory store update
+
+
+@dataclass(frozen=True)
+class Label:
+    kind: LabelKind
+    name: str = ""
+    input: Any = None
+    output: Any = None
+
+    def __repr__(self):
+        if self.kind == LabelKind.TAU:
+            return "τ"
+        return f"{self.kind.name.lower()}({self.name})"
+
+
+TAU = Label(LabelKind.TAU)
+
+
+# ============================================================
+# Continuation Frames (K)
+# ============================================================
+
+class Kont:
+    """Base class for continuation frames."""
+    pass
+
+
+class HaltK(Kont):
+    """Top-level: machine is done."""
+    def __repr__(self):
+        return "halt"
+
+
+class CompK(Kont):
+    """Composition: apply g to result of f."""
+    def __init__(self, g_term, prev: Kont):
+        self.g_term = g_term
+        self.prev = prev
+
+    def __repr__(self):
+        return f"compK({self.g_term}) :: {self.prev}"
+
+
+class LoopK(Kont):
+    """Loop: check condition after body evaluation."""
+    def __init__(self, body, cond, remaining: int, prev: Kont):
+        self.body = body
+        self.cond = cond
+        self.remaining = remaining
+        self.prev = prev
+
+    def __repr__(self):
+        return f"loopK(n={self.remaining}) :: {self.prev}"
+
+
+class PairLK(Kont):
+    """Pair: left branch done, need to start right."""
+    def __init__(self, g_term, input_val, prev: Kont):
+        self.g_term = g_term
+        self.input_val = input_val
+        self.prev = prev
+
+    def __repr__(self):
+        return f"pairLK({self.g_term}) :: {self.prev}"
+
+
+class PairRK(Kont):
+    """Pair: right branch evaluating, left value stored."""
+    def __init__(self, left_val, prev: Kont):
+        self.left_val = left_val
+        self.prev = prev
+
+    def __repr__(self):
+        return f"pairRK :: {self.prev}"
+
+
+class GuardK(Kont):
+    """Guard: check predicate on result."""
+    def __init__(self, predicate, retries: int, agent_term, input_val, prev: Kont):
+        self.predicate = predicate
+        self.retries = retries
+        self.agent_term = agent_term
+        self.input_val = input_val
+        self.prev = prev
+
+    def __repr__(self):
+        return f"guardK(k={self.retries}) :: {self.prev}"
+
+
+class MemK(Kont):
+    """Memory: store result after inner computation."""
+    def __init__(self, store_key: str, prev: Kont):
+        self.store_key = store_key
+        self.prev = prev
+
+    def __repr__(self):
+        return f"memK({self.store_key}) :: {self.prev}"
+
+
+# ============================================================
+# CEK Machine State
+# ============================================================
+
+@dataclass
+class CEKState:
+    """
+    Machine state = <C, E, K, σ, c>
+
+    C: control (current term — a lambdagent Term or a value)
+    E: environment (dict of bindings)
+    K: continuation (stack of frames)
+    store: memory store (dict)
+    cost: cumulative cost vector
+    """
+    control: Any           # Term or value
+    env: Dict[str, Any]    # environment
+    kont: Kont             # continuation stack
+    store: Dict[str, Any]  # memory store
+    cost: CostVector       # cumulative cost
+
+    def is_terminal(self) -> bool:
+        """Terminal if control is a value and continuation is halt."""
+        return isinstance(self.kont, HaltK) and self._is_value(self.control)
+
+    @staticmethod
+    def _is_value(x) -> bool:
+        """Values are non-Term Python objects (str, int, dict, tuple, etc.)
+        Auxiliary CEK terms (_AppTerm, _FstTerm, _SndTerm, _LoopState) are NOT values."""
+        from .core import Term
+        if isinstance(x, (Term, _AppTerm, _FstTerm, _SndTerm, _LoopState)):
+            return False
+        return True
+
+
+# ============================================================
+# Transition Trace Entry
+# ============================================================
+
+@dataclass
+class Transition:
+    """Record of one CEK transition."""
+    step: int
+    rule: str               # e.g. "C-Comp", "C-Lam", "C-CompRet"
+    label: Label
+    state_before: str       # string repr of state before
+    state_after: str        # string repr of state after
+    cost_delta: CostVector  # cost added in this step
+    duration_ms: float      # wall-clock time
+
+
+# ============================================================
+# Agent CEK Machine
+# ============================================================
+
+class AgentCEKMachine:
+    """
+    Agent CEK Machine — step-by-step executor for lambdagent programs.
+
+    Usage:
+        from lambdagent.cek_machine import AgentCEKMachine
+
+        machine = AgentCEKMachine()
+        result = machine.run(agent, input_val)
+
+        # Or step-by-step:
+        machine.load(agent, input_val)
+        while not machine.state.is_terminal():
+            transition = machine.step()
+            print(transition.rule, transition.label)
+        result = machine.state.control
+    """
+
+    def __init__(self, store: Dict[str, Any] | None = None):
+        self.state: Optional[CEKState] = None
+        self.trace: List[Transition] = []
+        self.step_count: int = 0
+        self._initial_store = store or {}
+
+    def load(self, term, input_val: Any) -> CEKState:
+        """Load a term into the machine, creating the initial state.
+
+        The initial control is the application (term input_val).
+        """
+        from .core import Term
+
+        # Create application term: term applied to input_val
+        app = _AppTerm(term, input_val)
+
+        self.state = CEKState(
+            control=app,
+            env={},
+            kont=HaltK(),
+            store=dict(self._initial_store),
+            cost=ZERO_COST,
+        )
+        self.trace = []
+        self.step_count = 0
+        return self.state
+
+    def _is_stuck(self) -> bool:
+        """Check if machine is stuck (control is non-value, non-dispatchable)."""
+        s = self.state
+        if s is None:
+            return True
+        ctrl = s.control
+        if s._is_value(ctrl):
+            return False
+        if isinstance(ctrl, (_AppTerm, _FstTerm, _SndTerm)):
+            return False
+        return True
+
+    def step(self) -> Transition:
+        """Execute one machine transition. Returns the transition record."""
+        if self.state is None:
+            raise RuntimeError("Machine not loaded. Call load() first.")
+        if self.state.is_terminal():
+            raise RuntimeError("Machine already halted.")
+
+        t0 = time.time()
+        state_before = self._state_repr()
+        cost_before = self.state.cost
+
+        rule, label = self._dispatch()
+
+        duration_ms = (time.time() - t0) * 1000
+        cost_delta = CostVector(
+            self.state.cost.tokens - cost_before.tokens,
+            self.state.cost.latency - cost_before.latency,
+            self.state.cost.money - cost_before.money,
+        )
+
+        self.step_count += 1
+        transition = Transition(
+            step=self.step_count,
+            rule=rule,
+            label=label,
+            state_before=state_before,
+            state_after=self._state_repr(),
+            cost_delta=cost_delta,
+            duration_ms=duration_ms,
+        )
+        self.trace.append(transition)
+        return transition
+
+    def run(self, term, input_val: Any, max_steps: int = 10000) -> Any:
+        """Run to completion. Returns the final value."""
+        self.load(term, input_val)
+        while not self.state.is_terminal():
+            if self.step_count >= max_steps:
+                raise RuntimeError(f"CEK machine exceeded {max_steps} steps")
+            self.step()
+        return self.state.control
+
+    def print_trace(self):
+        """Print the full transition trace."""
+        for t in self.trace:
+            cost_str = f" +{t.cost_delta}" if t.cost_delta.tokens > 0 else ""
+            print(f"  [{t.step:3d}] {t.rule:<16s} {t.label}{cost_str}  ({t.duration_ms:.1f}ms)")
+
+    # ── Internal dispatch ─────────────────────────────────────
+
+    def _dispatch(self) -> Tuple[str, Label]:
+        """Main dispatch: examine control and continuation, apply transition."""
+        from .primitives import Lam, Compose, If, Loop, Pair, Tool
+        from .extensions import Route, Guard, Memory, Par
+        from .core import Term
+
+        s = self.state
+        ctrl = s.control
+
+        # ── If control is a value, pop continuation ──
+        if s._is_value(ctrl):
+            return self._return(ctrl)
+
+        # ── Application: term applied to value ──
+        if isinstance(ctrl, _AppTerm):
+            return self._dispatch_app(ctrl.operator, ctrl.operand)
+
+        # ── Fst / Snd ──
+        if isinstance(ctrl, _FstTerm):
+            if isinstance(ctrl.inner, tuple) and len(ctrl.inner) == 2:
+                s.control = ctrl.inner[0]
+                return "C-Fst", TAU
+            # inner not yet a pair value — should not happen in well-formed programs
+            raise RuntimeError(f"Fst applied to non-pair: {ctrl.inner}")
+
+        if isinstance(ctrl, _SndTerm):
+            if isinstance(ctrl.inner, tuple) and len(ctrl.inner) == 2:
+                s.control = ctrl.inner[1]
+                return "C-Snd", TAU
+            raise RuntimeError(f"Snd applied to non-pair: {ctrl.inner}")
+
+        # ── Term that needs to be applied but has no operand ──
+        # This shouldn't happen in well-formed programs
+        raise RuntimeError(f"CEK stuck: control is {type(ctrl).__name__}, not a value or application")
+
+    def _dispatch_app(self, op, arg) -> Tuple[str, Label]:
+        """Dispatch an application (op arg)."""
+        from .primitives import Lam, Compose, If, Loop, Pair, Tool
+        from .extensions import Route, Guard, Memory, Par
+        from .core import Term, Context
+
+        s = self.state
+
+        # ── C-Lam: YIELD to LLM oracle ──
+        if isinstance(op, Lam):
+            t0 = time.time()
+            ctx = Context(bindings=s.env, memory=s.store)
+            result = op.apply(arg, ctx)
+            elapsed = time.time() - t0
+
+            # Extract cost from context trace
+            tokens = ctx.trace[-1].tokens_used if ctx.trace else 0
+            cost_llm = CostVector(
+                tokens=tokens,
+                latency=elapsed,
+                money=tokens * _price_per_token(op.model),
+            )
+            label = Label(LabelKind.LLM, op._name, arg, result)
+            s.control = result
+            s.cost = s.cost + cost_llm
+            return "C-Lam", label
+
+        # ── C-Tool: YIELD to tool oracle ──
+        if isinstance(op, Tool):
+            t0 = time.time()
+            ctx = Context(bindings=s.env, memory=s.store)
+            result = op.apply(arg, ctx)
+            elapsed = time.time() - t0
+
+            cost_tool = CostVector(latency=elapsed)
+            label = Label(LabelKind.TOOL, op._name, arg, result)
+            s.control = result
+            s.cost = s.cost + cost_tool
+            return "C-Tool", label
+
+        # ── C-Comp: push compK, evaluate f(v) ──
+        if isinstance(op, Compose):
+            stages = op.stages if hasattr(op, 'stages') else [op.f, op.g]
+            if len(stages) == 2:
+                f, g = stages
+                s.control = _AppTerm(f, arg)
+                s.kont = CompK(g, s.kont)
+                return "C-Comp", TAU
+            else:
+                # Multi-stage: (f >> g >> h)(v) → push all but first
+                first = stages[0]
+                rest = stages[1:]
+                # Build nested CompK: compK(h, compK(g, ...))
+                for stage in reversed(rest):
+                    s.kont = CompK(stage, s.kont)
+                s.control = _AppTerm(first, arg)
+                return "C-Comp", TAU
+
+        # ── C-If: dispatch on predicate ──
+        if isinstance(op, If):
+            cond = op.cond
+            if isinstance(cond, Term):
+                ctx = Context(bindings=s.env, memory=s.store)
+                cond_result = cond.apply(arg, ctx)
+                cond_result = If._is_truthy(cond_result)
+            else:
+                cond_result = cond(arg)
+            if cond_result:
+                s.control = _AppTerm(op.then_, arg)
+            else:
+                s.control = _AppTerm(op.else_, arg)
+            return "C-If", TAU
+
+        # ── C-Route: dispatch on classifier ──
+        if isinstance(op, Route):
+            label_val = op.classifier(arg) if callable(op.classifier) else op.classifier.apply(arg, Context())
+            if label_val in op.routes:
+                s.control = _AppTerm(op.routes[label_val], arg)
+            else:
+                raise RuntimeError(f"Route: no branch for label '{label_val}'")
+            return "C-Route", TAU
+
+        # ── C-Loop: check condition or unfold ──
+        if isinstance(op, (Loop, _LoopState)):
+            cond = op.condition
+            remaining = op.max_steps
+
+            if remaining == 0:
+                # E-LoopBound: return value
+                s.control = arg
+                return "C-LoopBound", TAU
+
+            # Loop.condition takes (result, step) — we pass step=0 here
+            # since the CEK machine tracks iteration via remaining count
+            step_idx = (op.max_steps if isinstance(op, Loop) else remaining) - remaining
+            try:
+                cond_result = cond(arg, step_idx)
+            except TypeError:
+                # Some conditions only take one arg
+                cond_result = cond(arg)
+
+            if cond_result:
+                # E-LoopBase: terminate
+                s.control = arg
+                return "C-LoopBase", TAU
+            else:
+                # E-LoopUnfold: evaluate body, push loopK
+                body = op.body
+                new_remaining = remaining - 1
+                s.control = _AppTerm(body, arg)
+                s.kont = LoopK(body, cond, new_remaining, s.kont)
+                return "C-LoopUnfold", TAU
+
+        # ── C-Pair: evaluate left branch, push pairLK ──
+        if isinstance(op, (Pair, Par)):
+            f = op.first if hasattr(op, 'first') else op.f
+            g = op.second if hasattr(op, 'second') else op.g
+            s.control = _AppTerm(f, arg)
+            s.kont = PairLK(g, arg, s.kont)
+            return "C-Pair", TAU
+
+        # ── C-Guard: evaluate inner agent, push guardK ──
+        if isinstance(op, Guard):
+            s.control = _AppTerm(op.agent, arg)
+            retries = op.retry
+            s.kont = GuardK(op.validator, retries, op.agent, arg, s.kont)
+            return "C-Guard", TAU
+
+        # ── C-Mem: evaluate inner agent, push memK ──
+        if isinstance(op, Memory):
+            s.control = _AppTerm(op.agent, arg)
+            store_key = op._name  # use Memory's name as store key
+            s.kont = MemK(store_key, s.kont)
+            return "C-Mem", TAU
+
+        # ── Fallback: try direct application ──
+        if isinstance(op, Term):
+            ctx = Context(bindings=s.env, memory=s.store)
+            result = op.apply(arg, ctx)
+            s.control = result
+            # Copy any trace/memory updates
+            s.store.update(ctx.memory)
+            return "C-Generic", TAU
+
+        raise RuntimeError(f"CEK dispatch: unknown operator type {type(op).__name__}")
+
+    def _return(self, val) -> Tuple[str, Label]:
+        """Return a value to the continuation."""
+        s = self.state
+        k = s.kont
+
+        # ── C-Halt ──
+        if isinstance(k, HaltK):
+            # Already terminal — should not be called
+            return "C-Halt", TAU
+
+        # ── C-CompRet: pass value to next stage ──
+        if isinstance(k, CompK):
+            s.control = _AppTerm(k.g_term, val)
+            s.kont = k.prev
+            return "C-CompRet", TAU
+
+        # ── C-LoopRet: feed body result back into loop ──
+        if isinstance(k, LoopK):
+            # Re-enter loop with the body's result as new value
+            new_loop = _LoopState(k.body, k.cond, k.remaining)
+            s.control = _AppTerm(new_loop, val)
+            s.kont = k.prev
+            return "C-LoopRet", TAU
+
+        # ── C-PairMid: left done, start right ──
+        if isinstance(k, PairLK):
+            s.control = _AppTerm(k.g_term, k.input_val)
+            s.kont = PairRK(val, k.prev)
+            return "C-PairMid", TAU
+
+        # ── C-PairRet: both done, construct pair ──
+        if isinstance(k, PairRK):
+            s.control = (k.left_val, val)
+            s.kont = k.prev
+            return "C-PairRet", TAU
+
+        # ── C-GuardRet: check predicate ──
+        if isinstance(k, GuardK):
+            try:
+                passed = k.predicate(val)
+            except Exception:
+                passed = False
+
+            if passed:
+                s.control = val
+                s.kont = k.prev
+                return "C-GuardOK", TAU
+            elif k.retries > 0:
+                # Retry: re-evaluate inner agent
+                s.control = _AppTerm(k.agent_term, k.input_val)
+                s.kont = GuardK(k.predicate, k.retries - 1, k.agent_term, k.input_val, k.prev)
+                return "C-GuardRetry", TAU
+            else:
+                s.control = None  # unit
+                s.kont = k.prev
+                return "C-GuardFail", TAU
+
+        # ── C-MemRet: store result ──
+        if isinstance(k, MemK):
+            s.store[k.store_key] = val
+            label = Label(LabelKind.MEM, k.store_key, None, val)
+            s.control = val
+            s.kont = k.prev
+            return "C-MemRet", label
+
+        raise RuntimeError(f"CEK return: unknown continuation {type(k).__name__}")
+
+    def _state_repr(self) -> str:
+        """Compact string representation of current state."""
+        s = self.state
+        ctrl_str = _term_repr(s.control)
+        return f"<{ctrl_str}, K={s.kont}, σ={len(s.store)} keys, {s.cost}>"
+
+
+# ============================================================
+# Internal helper terms (not part of source syntax)
+# ============================================================
+
+class _AppTerm:
+    """Application: operator applied to operand. Auxiliary term for CEK."""
+    def __init__(self, operator, operand):
+        self.operator = operator
+        self.operand = operand
+
+    def __repr__(self):
+        return f"({_term_repr(self.operator)} {_term_repr(self.operand)})"
+
+
+class _FstTerm:
+    """Fst projection. Auxiliary term for CEK."""
+    def __init__(self, inner):
+        self.inner = inner
+
+
+class _SndTerm:
+    """Snd projection. Auxiliary term for CEK."""
+    def __init__(self, inner):
+        self.inner = inner
+
+
+class _LoopState:
+    """Represents a Loop with reduced remaining count. Auxiliary for CEK."""
+    def __init__(self, body, cond, remaining: int):
+        self.body = body
+        self.condition = cond
+        self.max_steps = remaining
+
+
+def _term_repr(t) -> str:
+    """Compact representation of a term or value."""
+    if isinstance(t, _AppTerm):
+        return f"({_term_repr(t.operator)} {_term_repr(t.operand)})"
+    if isinstance(t, str):
+        s = t[:40] + "..." if len(t) > 40 else t
+        return f'"{s}"'
+    if isinstance(t, tuple):
+        return f"({_term_repr(t[0])}, {_term_repr(t[1])})"
+    if hasattr(t, '_name'):
+        return t._name
+    return repr(t)[:40]
+
+
+def _price_per_token(model: str) -> float:
+    """Rough per-token price for cost tracking."""
+    prices = {
+        "claude-sonnet-4-20250514": 3e-6,
+        "claude-opus-4-20250514": 15e-6,
+        "claude-haiku-4-5-20251001": 0.25e-6,
+        "gpt-4": 30e-6,
+        "gpt-4o": 2.5e-6,
+        "qwen3-max": 1e-6,
+    }
+    for key, price in prices.items():
+        if key in model:
+            return price
+    return 1e-6  # default