Add CLI with run, speculate, best-of-n, reflexion, and status commands

Signed-off-by: Cong Wang <cwang@multikernel.io>

Author: congwang-mk

pyproject.toml

@@ -9,6 +9,9 @@ description = "Unified branching for speculative execution: filesystem, process,
 requires-python = ">=3.10"
 license = "Apache-2.0"
 
+[project.scripts]
+branching = "cli:main"
+
 [tool.setuptools.packages.find]
 where = ["src"]
 

setup.cfg

@@ -11,3 +11,7 @@ packages = find:
 
 [options.packages.find]
 where = src
+
+[options.entry_points]
+console_scripts =
+    branching = cli:main

src/cli/__init__.py

@@ -0,0 +1,212 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CLI for BranchContext — run commands in COW branches from the terminal."""
+
+import argparse
+import json
+import sys
+import uuid
+from pathlib import Path
+from typing import Optional
+
+from branching.fs._mount import parse_mounts
+
+
+def _find_branchfs_mounts():
+    """Find all branchfs mounts (handles both 'fuse.branchfs' and bare 'fuse' types)."""
+    results = []
+    for m in parse_mounts():
+        if m.fstype == "fuse.branchfs":
+            results.append(m)
+        elif m.fstype == "fuse" and m.source == "branchfs":
+            results.append(m)
+    return results
+
+
+def _find_workspace_from_cwd() -> Optional[Path]:
+    """Find a branchfs mount that contains the current working directory.
+
+    Checks all branchfs mounts, returns the longest matching mountpoint
+    (most specific ancestor of cwd).
+    """
+    cwd = Path.cwd().resolve()
+    mounts = _find_branchfs_mounts()
+    best: Optional[Path] = None
+    for m in mounts:
+        mp = m.mountpoint.resolve()
+        try:
+            cwd.relative_to(mp)
+        except ValueError:
+            continue
+        if best is None or len(mp.parts) > len(best.parts):
+            best = mp
+    return best
+
+
+def _resolve_workspace(args: argparse.Namespace) -> Path:
+    """Resolve workspace path from --workspace arg or cwd auto-detection."""
+    if hasattr(args, "workspace") and args.workspace:
+        p = Path(args.workspace).resolve()
+        if not p.is_dir():
+            _print_error(f"workspace path does not exist: {p}", args)
+            sys.exit(1)
+        return p
+    detected = _find_workspace_from_cwd()
+    if detected is None:
+        _print_error(
+            "no branchfs mount found for current directory; use -w/--workspace",
+            args,
+        )
+        sys.exit(1)
+    return detected
+
+
+def _generate_branch_name(prefix: str = "run") -> str:
+    """Generate a branch name like 'run-a1b2c3d4'."""
+    return f"{prefix}-{uuid.uuid4().hex[:8]}"
+
+
+def _print_result(data: dict, args: argparse.Namespace) -> None:
+    """Print result as JSON (if --json) or human-readable text."""
+    if getattr(args, "json", False):
+        print(json.dumps(data))
+    else:
+        for key, value in data.items():
+            print(f"{key}: {value}")
+
+
+def _print_error(message: str, args: argparse.Namespace) -> None:
+    """Print error as JSON (if --json) or plain text to stderr."""
+    if getattr(args, "json", False):
+        print(json.dumps({"error": message}), file=sys.stderr)
+    else:
+        print(f"error: {message}", file=sys.stderr)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="branching",
+        description="CLI for BranchContext — COW branching for AI agent workflows.",
+    )
+    sub = parser.add_subparsers(dest="command")
+
+    # --- run ---
+    p_run = sub.add_parser(
+        "run",
+        help="Run a command in a new branch.",
+        description=(
+            "Run CMD in a new branch. Commits on exit 0, aborts on non-zero."
+        ),
+    )
+    p_run.add_argument("-w", "--workspace", help="Workspace path (auto-detected from cwd)")
+    p_run.add_argument("-b", "--branch", help="Branch name (auto-generated if omitted)")
+    p_run.add_argument(
+        "--on-error",
+        choices=["abort", "none"],
+        default="abort",
+        help="Action on non-zero exit (default: abort)",
+    )
+    p_run.add_argument(
+        "--ask",
+        action="store_true",
+        help="Prompt to commit/abort instead of auto-deciding",
+    )
+    p_run.add_argument("--json", action="store_true", help="JSON output")
+    p_run.add_argument("cmd", nargs=argparse.REMAINDER, metavar="CMD",
+                        help="Command to run (after --)")
+
+    # --- speculate ---
+    p_spec = sub.add_parser(
+        "speculate",
+        help="Race N commands in parallel branches.",
+        description="Race N commands in parallel; first success wins.",
+    )
+    p_spec.add_argument("-w", "--workspace", help="Workspace path (auto-detected from cwd)")
+    p_spec.add_argument(
+        "-c", "--candidate",
+        action="append",
+        required=True,
+        dest="candidates",
+        metavar="CMD",
+        help="Command string to run as candidate (use multiple -c flags)",
+    )
+    p_spec.add_argument("--timeout", type=float, default=None, help="Timeout in seconds")
+    p_spec.add_argument("--json", action="store_true", help="JSON output")
+
+    # --- best-of-n ---
+    p_bon = sub.add_parser(
+        "best-of-n",
+        help="Run CMD N times in parallel, commit the highest-scoring success.",
+        description="Run CMD N times in parallel branches, commit the highest-scoring success.",
+    )
+    p_bon.add_argument("-w", "--workspace", help="Workspace path (auto-detected from cwd)")
+    p_bon.add_argument("-n", type=int, default=3, help="Number of parallel attempts (default: 3)")
+    p_bon.add_argument("--timeout", type=float, default=None, help="Timeout in seconds")
+    p_bon.add_argument("--json", action="store_true", help="JSON output")
+    p_bon.add_argument("cmd", nargs=argparse.REMAINDER, metavar="CMD",
+                        help="Command to run (after --)")
+
+    # --- reflexion ---
+    p_refl = sub.add_parser(
+        "reflexion",
+        help="Sequential retry with optional critique feedback loop.",
+        description="Sequential retry with optional critique feedback loop.",
+    )
+    p_refl.add_argument("-w", "--workspace", help="Workspace path (auto-detected from cwd)")
+    p_refl.add_argument("--retries", type=int, default=3, help="Max retry attempts (default: 3)")
+    p_refl.add_argument("--critique", type=str, default=None,
+                         help="Shell command to generate feedback after failure")
+    p_refl.add_argument("--json", action="store_true", help="JSON output")
+    p_refl.add_argument("cmd", nargs=argparse.REMAINDER, metavar="CMD",
+                         help="Command to run (after --)")
+
+    # --- status ---
+    p_status = sub.add_parser(
+        "status",
+        help="Show workspace info and branch list.",
+    )
+    p_status.add_argument("-w", "--workspace", help="Workspace path (auto-detected from cwd)")
+    p_status.add_argument("--json", action="store_true", help="JSON output")
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    if args.command is None:
+        parser.print_help()
+        sys.exit(0)
+
+    # Strip leading '--' from remainder args for commands that take CMD...
+    if args.command in ("run", "best-of-n", "reflexion"):
+        cmd = args.cmd
+        if cmd and cmd[0] == "--":
+            cmd = cmd[1:]
+        args.cmd = cmd
+
+    try:
+        if args.command == "run":
+            from .run import cmd_run
+            sys.exit(cmd_run(args))
+        elif args.command == "speculate":
+            from .speculate import cmd_speculate
+            sys.exit(cmd_speculate(args))
+        elif args.command == "best-of-n":
+            from .best_of_n import cmd_best_of_n
+            sys.exit(cmd_best_of_n(args))
+        elif args.command == "reflexion":
+            from .reflexion import cmd_reflexion
+            sys.exit(cmd_reflexion(args))
+        elif args.command == "status":
+            from .status import cmd_status
+            sys.exit(cmd_status(args))
+        else:
+            parser.print_help()
+            sys.exit(1)
+    except KeyboardInterrupt:
+        print(file=sys.stderr)
+        sys.exit(130)
+    except Exception as e:
+        _print_error(str(e), args)
+        sys.exit(1)

src/cli/best_of_n.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Implementation of 'branching best-of-n' command."""
+
+import json
+import os
+import subprocess
+from pathlib import Path
+
+from branching import BestOfN, Workspace
+
+from . import _print_error, _resolve_workspace
+
+
+def _make_task(cmd: list[str]):
+    """Wrap a command into a BestOfN task callable.
+
+    Returns a callable(path, index) -> (success, score).
+    The child process can write a score float to fd 3.
+    """
+
+    def task(workdir: Path, index: int) -> tuple[bool, float]:
+        # Create a pipe for the child to report its score.
+        # Python 3.4+ creates pipe fds with CLOEXEC, so they are
+        # automatically closed on exec — only fd 3 (dup2 clears
+        # CLOEXEC) survives into the child.
+        r_fd, w_fd = os.pipe()
+
+        env = {**os.environ, "BRANCHING_ATTEMPT": str(index)}
+
+        def _preexec():
+            os.dup2(w_fd, 3)
+
+        proc = subprocess.Popen(
+            cmd,
+            cwd=workdir,
+            env=env,
+            close_fds=False,
+            preexec_fn=_preexec,
+        )
+        # Close write end in the parent so read will EOF when child exits
+        os.close(w_fd)
+
+        proc.wait()
+        success = proc.returncode == 0
+
+        # Read score from the pipe
+        with os.fdopen(r_fd, "r") as f:
+            raw = f.read().strip()
+
+        if raw:
+            try:
+                score = float(raw)
+            except ValueError:
+                score = 1.0 if success else 0.0
+        else:
+            score = 1.0 if success else 0.0
+
+        return (success, score)
+
+    return task
+
+
+def cmd_best_of_n(args) -> int:
+    if not args.cmd:
+        _print_error("no command specified; use -- CMD...", args)
+        return 1
+
+    ws_path = _resolve_workspace(args)
+    ws = Workspace(ws_path)
+
+    task = _make_task(args.cmd)
+    best = BestOfN(task, n=args.n, timeout=args.timeout)
+    outcome = best(ws)
+
+    results_summary = []
+    for r in outcome.all_results:
+        entry = {
+            "index": r.branch_index,
+            "success": r.success,
+            "score": r.score,
+        }
+        if r.exception is not None:
+            entry["error"] = str(r.exception)
+        results_summary.append(entry)
+
+    data = {
+        "command": "best-of-n",
+        "n": args.n,
+        "committed": outcome.committed,
+        "winner": {
+            "index": outcome.winner.branch_index,
+            "score": outcome.winner.score,
+        } if outcome.winner else None,
+        "results": results_summary,
+    }
+
+    if getattr(args, "json", False):
+        print(json.dumps(data))
+    else:
+        if outcome.committed:
+            print(
+                f"winner: attempt {outcome.winner.branch_index}"
+                f" (score: {outcome.winner.score:.2f})"
+            )
+        else:
+            print("no attempt succeeded")
+        for r in outcome.all_results:
+            status = "ok" if r.success else "fail"
+            print(f"  [{r.branch_index}] {status:<4} ({r.score:.2f})")
+
+    return 0 if outcome.committed else 1