levain-cli 0.1.0__py3-none-any.whl

levain_recipe/reducers.py

@@ -0,0 +1,66 @@
+r"""Blessed reducers for recipe state fields.
+
+LangGraph applies a reducer when a node returns an update for a
+state field.  The default merges by overwrite ("last write wins"),
+which is wrong for accumulating lists — messages, logs, any
+append-only channel.  Recipes opt into a reducer by name via
+``\* @reducer <name>`` above a variable declaration; validation
+rejects anything not in this registry.
+
+Keeping the set small and named (rather than accepting arbitrary
+dotted paths) keeps the DSL's surface area bounded: a customer
+recipe can't reach into the platform by declaring
+``@reducer some.module:func``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class ReducerSpec:
+    """Codegen info for a blessed reducer.
+
+    Deliberately metadata only — names and import lines, no
+    callables.  Binding a name to the function LangGraph invokes is
+    the runtime's job (``levain_sdk.reducers.RUNTIME_REDUCERS``),
+    and keeping it there is what lets this package validate a
+    recipe without installing the agent runtime.
+
+    ``element_type_name`` is the type each list item takes in the
+    generated ``_state.py``; ``import_line`` is the single import
+    codegen writes when this reducer is used.
+    """
+
+    name: str
+    element_type_name: str
+    import_line: str
+
+
+REDUCERS: dict[str, ReducerSpec] = {
+    "add_messages": ReducerSpec(
+        name="add_messages",
+        element_type_name="Message",
+        import_line="from levain_sdk.messages import Message, add_messages",
+    ),
+    "add_str": ReducerSpec(
+        name="add_str",
+        element_type_name="str",
+        import_line="from levain_sdk.reducers import add_str",
+    ),
+    "extend_list": ReducerSpec(
+        name="extend_list",
+        # ``dict[str, Any]`` rather than bare ``dict``: the latter
+        # trips mypy's ``type-arg`` rule on the generated
+        # ``_state.py``.
+        element_type_name="dict[str, Any]",
+        import_line="from levain_sdk.reducers import extend_list",
+    ),
+}
+
+if not all(k == v.name for k, v in REDUCERS.items()):
+    raise ValueError("REDUCERS dict key must equal ReducerSpec.name")
+
+
+__all__ = ["REDUCERS", "ReducerSpec"]

levain_recipe/source_check.py

@@ -0,0 +1,279 @@
+"""Python source validation for recipe packages.
+
+Checks that @node function references in graph.tla resolve
+to actual Python functions in the recipe's src/ directory.
+Uses ast (no imports) for safe static analysis.
+"""
+
+from __future__ import annotations
+
+import ast
+import os
+import re
+import tomllib
+from pathlib import Path
+
+from .models import NodeAnnotation, ValidatedSpec, ValidationError
+
+_SDK_PACKAGE = "levain-sdk"
+_NORMALIZE_RE = re.compile(r"[-_.]+")
+
+
+def check_python_sources(
+    spec: ValidatedSpec,
+    recipe_dir: str,
+) -> list[ValidationError]:
+    """Validate Python source references in a recipe package.
+
+    Checks:
+    1. module:function refs resolve to files + functions
+    2. Functions are async def with 2 parameters
+    3. pyproject.toml exists and is valid
+    4. pyproject.toml declares SDK dependency
+    5. No orphan public ``async def`` functions in node-hosting
+       files — every top-level async function in a file referenced
+       by a ``@node`` annotation must itself be bound to a
+       ``@node``, or be marked private with a leading underscore.
+       Catches recipes that add Python functions for new behaviour
+       but forget to wire them into ``graph.tla``.
+    """
+    errors: list[ValidationError] = []
+    src_dir = os.path.join(recipe_dir, "src")
+
+    errors.extend(_check_pyproject(recipe_dir))
+    errors.extend(_check_function_refs(spec.nodes, src_dir))
+    errors.extend(_check_orphan_node_functions(spec.nodes, src_dir))
+
+    return errors
+
+
+def _check_pyproject(recipe_dir: str) -> list[ValidationError]:
+    """Check pyproject.toml exists and declares SDK dependency."""
+    path = os.path.join(recipe_dir, "pyproject.toml")
+    if not os.path.isfile(path):
+        return [
+            ValidationError(
+                stage="source",
+                message="pyproject.toml not found in recipe directory.",
+            )
+        ]
+
+    try:
+        with open(path, "rb") as f:
+            data = tomllib.load(f)
+    except tomllib.TOMLDecodeError as exc:
+        return [
+            ValidationError(
+                stage="source",
+                message=f"pyproject.toml is not valid TOML: {exc}",
+            )
+        ]
+
+    deps = data.get("project", {}).get("dependencies", [])
+    has_sdk = any(
+        _NORMALIZE_RE.sub("-", re.split(r"[~><!=;\[ ]", d)[0].strip()).lower()
+        == _SDK_PACKAGE
+        for d in deps
+    )
+    if not has_sdk:
+        return [
+            ValidationError(
+                stage="source",
+                message=(
+                    f"pyproject.toml does not declare "
+                    f"'{_SDK_PACKAGE}' as a dependency."
+                ),
+            )
+        ]
+
+    return []
+
+
+def _check_function_refs(
+    nodes: list[NodeAnnotation],
+    src_dir: str,
+) -> list[ValidationError]:
+    """Check that module:function refs resolve to Python code."""
+    errors: list[ValidationError] = []
+
+    for node in nodes:
+        func = node.function
+
+        # Bare names (no colon) are dispatcher-handled node types
+        # (e.g. claude_agent, shell, recipe) — not Python refs.
+        if ":" not in func:
+            continue
+
+        module_path, func_name = func.rsplit(":", 1)
+
+        # Resolve module path to file.
+        parts = module_path.split(".")
+        py_file = Path(src_dir, *parts).with_suffix(".py")
+
+        # Guard against path traversal escaping src/.
+        try:
+            py_file.resolve().relative_to(Path(src_dir).resolve())
+        except ValueError:
+            errors.append(
+                ValidationError(
+                    stage="source",
+                    message=(
+                        f"Module path '{module_path}' "
+                        f"escapes the src/ directory."
+                    ),
+                )
+            )
+            continue
+
+        if not py_file.is_file():
+            errors.append(
+                ValidationError(
+                    stage="source",
+                    message=(
+                        f"Module '{module_path}' not found "
+                        f"at {py_file.relative_to(Path(src_dir).parent)}."
+                    ),
+                )
+            )
+            continue
+
+        # Parse and check function exists.
+        try:
+            source = py_file.read_text(encoding="utf-8")
+            tree = ast.parse(source, filename=str(py_file))
+        except (SyntaxError, OSError) as exc:
+            errors.append(
+                ValidationError(
+                    stage="source",
+                    message=(f"Syntax error in {py_file.name}: {exc}"),
+                )
+            )
+            continue
+
+        func_node = _find_async_def(tree, func_name)
+        if func_node is None:
+            errors.append(
+                ValidationError(
+                    stage="source",
+                    message=(
+                        f"Function '{func_name}' not found as "
+                        f"a top-level async def in "
+                        f"{py_file.name}."
+                    ),
+                )
+            )
+            continue
+
+        # Check parameter count (state, config).
+        n_params = (
+            len(func_node.args.posonlyargs)
+            + len(func_node.args.args)
+            + len(func_node.args.kwonlyargs)
+        )
+        if n_params != 2:
+            errors.append(
+                ValidationError(
+                    stage="source",
+                    message=(
+                        f"Function '{func_name}' in "
+                        f"{py_file.name} has {n_params} "
+                        f"parameter(s), expected 2 "
+                        f"(state, config)."
+                    ),
+                )
+            )
+
+    return errors
+
+
+def _find_async_def(
+    tree: ast.Module,
+    name: str,
+) -> ast.AsyncFunctionDef | None:
+    """Find a top-level async def by name."""
+    for node in ast.iter_child_nodes(tree):
+        if isinstance(node, ast.AsyncFunctionDef) and node.name == name:
+            return node
+    return None
+
+
+def _check_orphan_node_functions(
+    nodes: list[NodeAnnotation],
+    src_dir: str,
+) -> list[ValidationError]:
+    """Flag public ``async def`` functions not bound to a ``@node``.
+
+    Recipes that add behaviour (new "nodes" in the agent sense)
+    must wire it through ``graph.tla`` — adding the function to
+    ``nodes.py`` alone leaves the recipe topology unchanged and
+    the new function as dead code.
+
+    The check is per-file: for each ``.py`` file referenced by
+    any ``@node`` annotation, collect the bound function names,
+    then flag any other top-level ``async def`` that doesn't
+    start with ``_`` (private convention).  Files not hosting
+    any node binding are ignored — utility modules can have
+    arbitrary public async helpers.
+
+    The error message stays neutral: it reports the
+    inconsistency and lets the recipe author / build agent
+    decide whether the right fix is to wire the function in or
+    to remove it as dead code.  We surface the gap; we don't
+    push a direction.
+    """
+    errors: list[ValidationError] = []
+
+    # file path -> set of bound function names from @node refs
+    bindings_per_file: dict[Path, set[str]] = {}
+    for node in nodes:
+        func = node.function
+        if ":" not in func:
+            continue
+        module_path, func_name = func.rsplit(":", 1)
+        py_file = Path(src_dir, *module_path.split(".")).with_suffix(".py")
+        try:
+            py_file.resolve().relative_to(Path(src_dir).resolve())
+        except ValueError:
+            continue
+        if not py_file.is_file():
+            continue
+        bindings_per_file.setdefault(py_file, set()).add(func_name)
+
+    for py_file, bound_names in bindings_per_file.items():
+        try:
+            tree = ast.parse(py_file.read_text(encoding="utf-8"))
+        except (SyntaxError, OSError):
+            # ``_check_function_refs`` already surfaced this; don't
+            # double-report.
+            continue
+        async_defs = {
+            n.name
+            for n in ast.iter_child_nodes(tree)
+            if isinstance(n, ast.AsyncFunctionDef)
+        }
+        # If any bound name is missing from the file,
+        # _check_function_refs already reported it — skip
+        # the orphan scan to avoid double-reporting every
+        # other function in the file as unbound.
+        if not bound_names.issubset(async_defs):
+            continue
+        for child in ast.iter_child_nodes(tree):
+            if not isinstance(child, ast.AsyncFunctionDef):
+                continue
+            if child.name.startswith("_"):
+                continue
+            if child.name in bound_names:
+                continue
+            errors.append(
+                ValidationError(
+                    stage="source",
+                    message=(
+                        f"Async function '{child.name}' in "
+                        f"{py_file.name} isn't bound to a @node "
+                        f"annotation in graph.tla. Either wire it "
+                        f"into the graph or remove it as dead code."
+                    ),
+                )
+            )
+
+    return errors

levain_recipe/tlc.py

@@ -0,0 +1,236 @@
+"""Stage 5: TLC model checking via tla2tools.jar."""
+
+from __future__ import annotations
+
+import os
+import re
+import subprocess
+import tempfile
+from typing import Any
+
+from ._java import RE_MODULE_NAME, require_java, resolve_jar
+from .models import ValidationError
+
+_TLC_TIMEOUT_S = 60
+
+
+def model_check(
+    translated_source: str,
+    module_name: str,
+    invariants: list[str],
+) -> list[ValidationError] | None:
+    """Run TLC on the translated TLA+ spec.
+
+    Parameters
+    ----------
+    translated_source
+        TLA+ source after PlusCal translation.
+    module_name
+        Module name for temp file naming.
+    invariants
+        Invariant names to check.
+
+    Returns
+    -------
+    None | list[ValidationError]
+        None on success, errors on failure.
+    """
+    if not RE_MODULE_NAME.match(module_name):
+        return [
+            ValidationError(
+                stage="tlc",
+                message="Invalid module name.",
+            )
+        ]
+
+    jar = resolve_jar()
+    if jar is None:
+        env_jar = os.environ.get("TLA2TOOLS_JAR")
+        if env_jar:
+            msg = f"TLA2TOOLS_JAR={env_jar!r} does not exist."
+        else:
+            msg = "tla2tools.jar not found. Set TLA2TOOLS_JAR env var."
+        return [ValidationError(stage="tlc", message=msg)]
+
+    java_or_err = require_java("tlc")
+    if isinstance(java_or_err, list):
+        return java_or_err
+    java = java_or_err
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        spec_path = os.path.join(tmpdir, f"{module_name}.tla")
+        cfg_path = os.path.join(tmpdir, f"{module_name}.cfg")
+
+        with open(spec_path, "w", encoding="utf-8") as f:
+            f.write(translated_source)
+
+        cfg = _generate_cfg(invariants)
+        with open(cfg_path, "w", encoding="utf-8") as f:
+            f.write(cfg)
+
+        try:
+            result = subprocess.run(
+                [
+                    java,
+                    "-Xmx512m",
+                    "-jar",
+                    jar,
+                    "-config",
+                    cfg_path,
+                    "-deadlock",
+                    spec_path,
+                ],
+                capture_output=True,
+                text=True,
+                timeout=_TLC_TIMEOUT_S,
+                cwd=tmpdir,
+                check=False,
+            )
+        except subprocess.TimeoutExpired:
+            return [
+                ValidationError(
+                    stage="tlc",
+                    message=(
+                        "TLC model check timed out after "
+                        f"{_TLC_TIMEOUT_S}s. The state "
+                        "space may be too large."
+                    ),
+                )
+            ]
+
+        return _parse_tlc_output(
+            result.stdout + result.stderr,
+            result.returncode,
+        )
+
+
+def _generate_cfg(invariants: list[str]) -> str:
+    """Generate a TLC .cfg file.
+
+    Safety only — emits ``INVARIANT`` lines for each ``define``
+    block operator.  Liveness (``PROPERTY Termination``) is
+    intentionally NOT checked here: under PlusCal's restricted
+    profile, TLA+ can't cleanly express "the agent eventually
+    converges", and a mechanical bound (``retry_count``) would
+    limit agents arbitrarily.  Termination is instead enforced
+    at runtime by the orchestrator's budget/timeout guardrails.
+    """
+    lines = ["SPECIFICATION Spec", ""]
+
+    if invariants:
+        for inv in invariants:
+            lines.append(f"INVARIANT {inv}")
+        lines.append("")
+
+    return "\n".join(lines)
+
+
+# ── TLC output parsing ──
+
+_RE_ERROR = re.compile(r"Error:\s*(.*)")
+_RE_INVARIANT = re.compile(r"Invariant\s+(\w+)\s+is violated")
+_RE_STATE = re.compile(r"State\s+(\d+):\s*<(\w+)\s+line\s+(\d+)")
+_RE_VAR_VAL = re.compile(r"/\\\s*(\w+)\s*=\s*(.*)")
+
+
+def _parse_tlc_output(
+    output: str,
+    returncode: int,
+) -> list[ValidationError] | None:
+    """Parse TLC stdout/stderr for pass or failure."""
+    if returncode == 0 and "Model checking completed" in output:
+        return None
+
+    if returncode == 0 and "Model checking completed" not in output:
+        return [
+            ValidationError(
+                stage="tlc",
+                message=(
+                    "TLC exited successfully but did not report completion."
+                ),
+            )
+        ]
+
+    errors: list[ValidationError] = []
+
+    # Check for invariant violations.
+    trace = _extract_trace(output) if _RE_INVARIANT.search(output) else None
+    for inv_match in _RE_INVARIANT.finditer(output):
+        errors.append(
+            ValidationError(
+                stage="tlc",
+                message=(f"Invariant '{inv_match.group(1)}' is violated."),
+                counterexample=trace or None,
+            )
+        )
+    if errors:
+        return errors
+
+    # Generic errors.
+    for m in _RE_ERROR.finditer(output):
+        errors.append(
+            ValidationError(
+                stage="tlc",
+                message=m.group(1).strip(),
+            )
+        )
+
+    if not errors and returncode != 0:
+        errors.append(
+            ValidationError(
+                stage="tlc",
+                message=(
+                    f"TLC exited with code {returncode}: " + output[:200]
+                ),
+            )
+        )
+
+    return errors or None
+
+
+def _extract_trace(
+    output: str,
+) -> list[dict[str, Any]]:
+    """Extract the counterexample trace from TLC output."""
+    states: list[dict[str, Any]] = []
+    current_state: dict[str, Any] = {}
+
+    for line in output.splitlines():
+        state_m = _RE_STATE.match(line.strip())
+        if state_m:
+            if current_state:
+                states.append(current_state)
+            current_state = {
+                "state_num": int(state_m.group(1)),
+                "action": state_m.group(2),
+                "line": int(state_m.group(3)),
+                "vars": {},
+            }
+            continue
+
+        var_m = _RE_VAR_VAL.match(line.strip())
+        if var_m and current_state:
+            name = var_m.group(1)
+            value = _parse_tlc_value(var_m.group(2).strip())
+            current_state["vars"][name] = value
+
+    if current_state:
+        states.append(current_state)
+
+    return states
+
+
+def _parse_tlc_value(raw: str) -> Any:
+    """Best-effort parse of a TLC value."""
+    if raw == "TRUE":
+        return True
+    if raw == "FALSE":
+        return False
+    try:
+        return int(raw)
+    except ValueError:
+        pass
+    # Strip TLA+ string quotes.
+    if raw.startswith('"') and raw.endswith('"'):
+        return raw[1:-1]
+    return raw