levain-cli 0.1.0__py3-none-any.whl

levain_recipe/auth.py

@@ -0,0 +1,322 @@
+"""Device-flow login, token storage, and refresh for the Levain CLI.
+
+Authoring needs a credential for git, which can't do OAuth. Rather
+than have the user keep a long-lived API key on disk — or worse, hand
+one to an agent to paste into commands — the CLI signs in once with
+the OAuth **device authorization grant** and keeps a refresh token.
+Everything after that is derived and short-lived.
+
+Device flow rather than a loopback redirect on purpose: developers
+using a coding agent are often inside a container or an SSH session,
+where the browser is on a different machine and binding a local port
+proves nothing. Device flow only needs the user to open a URL
+wherever their browser happens to be.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import os
+import stat
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+from urllib.parse import urljoin
+
+import httpx
+
+DEFAULT_API = "https://api.levainlabs.com"
+DEFAULT_GIT_HOST = "git.levainlabs.com"
+
+# Refresh a little early: a token that expires mid-request is
+# indistinguishable from a revoked one to the caller.
+_EXPIRY_SKEW_S = 60
+
+_SCOPES = "openid profile email offline_access"
+
+
+class AuthError(Exception):
+    """Login is missing, expired, or the provider refused."""
+
+
+def api_base() -> str:
+    """API root to authenticate against."""
+    return os.environ.get("LEVAIN_API", DEFAULT_API).rstrip("/")
+
+
+def git_host() -> str:
+    """Only host the git credential helper may mint credentials for."""
+    return os.environ.get("LEVAIN_GIT_HOST", DEFAULT_GIT_HOST)
+
+
+def _store_path() -> Path:
+    root = os.environ.get("LEVAIN_HOME")
+    base = Path(root) if root else Path.home() / ".levain"
+    return base / "credentials.json"
+
+
+def _load_store() -> dict[str, Any]:
+    path = _store_path()
+    try:
+        with open(path, encoding="utf-8") as fh:
+            data = json.load(fh)
+    except (OSError, ValueError):
+        return {}
+    return data if isinstance(data, dict) else {}
+
+
+def _save_store(store: dict[str, Any]) -> None:
+    """Write the credential store, owner-readable only.
+
+    Created with restrictive permissions *before* anything is written,
+    so the secret is never briefly world-readable on a fresh install.
+    """
+    path = _store_path()
+    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
+    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+    with os.fdopen(fd, "w", encoding="utf-8") as fh:
+        json.dump(store, fh, indent=2)
+    os.chmod(path, stat.S_IRUSR | stat.S_IWUSR)
+
+
+@dataclass(frozen=True)
+class AuthConfig:
+    """Where to authenticate for a given deployment."""
+
+    authkit_domain: str
+    client_id: str
+    # Servers that know their identity provider serve the endpoints
+    # directly; empty means "fall back to OpenID discovery".
+    device_authorization_endpoint: str = ""
+    token_endpoint: str = ""
+
+
+def fetch_auth_config(api: str, client: httpx.Client) -> AuthConfig:
+    """Ask the deployment which identity provider to use.
+
+    Deliberately not baked into the package: a self-hosted install has
+    its own AuthKit tenant, and a client that hard-codes ours would
+    silently sign people into the wrong place.
+    """
+    resp = client.get(urljoin(api + "/", "api/v1/authoring/auth-config"))
+    if resp.status_code != 200:
+        raise AuthError(
+            f"{api} did not return a login configuration "
+            f"(HTTP {resp.status_code}). Check LEVAIN_API."
+        )
+    body = resp.json()
+    domain = str(body.get("authkit_domain", "")).rstrip("/")
+    client_id = str(body.get("client_id", ""))
+    if not domain or not client_id:
+        raise AuthError("This deployment has no CLI login configured.")
+    return AuthConfig(
+        authkit_domain=domain,
+        client_id=client_id,
+        device_authorization_endpoint=str(
+            body.get("device_authorization_endpoint", "") or ""
+        ),
+        token_endpoint=str(body.get("token_endpoint", "") or ""),
+    )
+
+
+def _discover(domain: str, client: httpx.Client) -> dict[str, Any]:
+    resp = client.get(f"{domain}/.well-known/openid-configuration")
+    if resp.status_code != 200:
+        raise AuthError(f"Could not read OpenID configuration from {domain}")
+    data: dict[str, Any] = resp.json()
+    return data
+
+
+def _flow_endpoints(
+    config: AuthConfig, client: httpx.Client, *, need_device: bool = True
+) -> tuple[str, str]:
+    """The (device authorization, token) endpoints to use.
+
+    Prefer what the deployment serves: WorkOS accepts AuthKit
+    application clients only on its native ``user_management``
+    endpoints, while the ones its discovery document advertises
+    refuse them. Discovery stays as the fallback for deployments
+    that serve no endpoints (any standards-compliant provider).
+
+    ``need_device`` is False for the refresh path, which only ever
+    uses the token endpoint: a standards-compliant provider need not
+    advertise ``device_authorization_endpoint`` in its discovery
+    document to support refresh.
+    """
+    if config.device_authorization_endpoint and config.token_endpoint:
+        return config.device_authorization_endpoint, config.token_endpoint
+    meta = _discover(config.authkit_domain, client)
+    device = meta.get("device_authorization_endpoint")
+    token = meta.get("token_endpoint")
+    if not token or (need_device and not device):
+        raise AuthError(
+            "This identity provider does not support device login."
+        )
+    return str(device or ""), str(token)
+
+
+@dataclass(frozen=True)
+class DevicePrompt:
+    """What the user has to do to finish signing in."""
+
+    verification_uri: str
+    user_code: str
+    interval: int
+    device_code: str
+    token_endpoint: str
+    client_id: str
+
+
+def start_device_login(api: str, client: httpx.Client) -> DevicePrompt:
+    """Begin the device flow; returns what to show the user."""
+    config = fetch_auth_config(api, client)
+    endpoint, token_endpoint = _flow_endpoints(config, client)
+    resp = client.post(
+        endpoint,
+        data={"client_id": config.client_id, "scope": _SCOPES},
+    )
+    if resp.status_code != 200:
+        raise AuthError(f"Could not start login (HTTP {resp.status_code}).")
+    body = resp.json()
+    device_code = body.get("device_code")
+    if not device_code:
+        raise AuthError(
+            "This identity provider returned an incomplete response."
+        )
+    return DevicePrompt(
+        verification_uri=str(
+            body.get("verification_uri_complete")
+            or body.get("verification_uri", "")
+        ),
+        user_code=str(body.get("user_code", "")),
+        interval=int(body.get("interval", 5) or 5),
+        device_code=str(device_code),
+        token_endpoint=token_endpoint,
+        client_id=config.client_id,
+    )
+
+
+def poll_device_login(
+    api: str,
+    prompt: DevicePrompt,
+    client: httpx.Client,
+    *,
+    timeout_s: int = 600,
+) -> None:
+    """Poll until the user approves, then store the tokens.
+
+    Honours the provider's back-off signals rather than hammering:
+    ``authorization_pending`` means keep waiting, ``slow_down`` means
+    wait longer, anything else is terminal.
+    """
+    deadline = time.time() + timeout_s
+    interval = prompt.interval
+    while time.time() < deadline:
+        time.sleep(interval)
+        resp = client.post(
+            prompt.token_endpoint,
+            data={
+                "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+                "device_code": prompt.device_code,
+                "client_id": prompt.client_id,
+            },
+        )
+        if resp.status_code == 200:
+            _store_tokens(api, resp.json())
+            return
+        error = ""
+        try:
+            error = str(resp.json().get("error", ""))
+        except ValueError:
+            pass
+        if error == "authorization_pending":
+            continue
+        if error == "slow_down":
+            interval += 5
+            continue
+        if error == "expired_token":
+            raise AuthError("The login request expired. Run it again.")
+        if error == "access_denied":
+            raise AuthError("Login was declined.")
+        raise AuthError(f"Login failed ({error or resp.status_code}).")
+    raise AuthError("Timed out waiting for approval.")
+
+
+def _jwt_exp(token: str) -> int:
+    """Best-effort ``exp`` claim of a JWT; 0 when unreadable."""
+    try:
+        payload = token.split(".")[1]
+        padded = payload + "=" * (-len(payload) % 4)
+        claims = json.loads(base64.urlsafe_b64decode(padded))
+        return int(claims.get("exp", 0))
+    except (ValueError, IndexError, TypeError, AttributeError):
+        return 0
+
+
+def _store_tokens(api: str, payload: dict[str, Any]) -> None:
+    access = str(payload.get("access_token", ""))
+    expires_in = payload.get("expires_in")
+    if expires_in is not None:
+        expires_at = int(time.time()) + int(expires_in)
+    else:
+        # WorkOS's native token responses carry no ``expires_in``;
+        # the JWT's own ``exp`` is authoritative. Unreadable means
+        # treat as already expired, so the next call refreshes.
+        expires_at = _jwt_exp(access)
+    store = _load_store()
+    store[api] = {
+        "access_token": access,
+        "refresh_token": payload.get("refresh_token", ""),
+        "expires_at": expires_at,
+    }
+    _save_store(store)
+
+
+def logout(api: str) -> bool:
+    """Forget the stored credential. True when there was one."""
+    store = _load_store()
+    if api not in store:
+        return False
+    del store[api]
+    _save_store(store)
+    return True
+
+
+def access_token(api: str, client: httpx.Client) -> str:
+    """A valid access token, refreshing if needed.
+
+    Raises :class:`AuthError` when the user has never logged in or the
+    refresh token is no longer accepted — both mean "run levain
+    login", which is what the caller tells them.
+    """
+    entry = _load_store().get(api)
+    if not entry:
+        raise AuthError(f"Not logged in to {api}. Run: levain login")
+    if int(entry.get("expires_at", 0)) - _EXPIRY_SKEW_S > time.time():
+        token: str = entry.get("access_token", "")
+        if token:
+            return token
+    refresh = entry.get("refresh_token", "")
+    if not refresh:
+        raise AuthError(f"Session expired for {api}. Run: levain login")
+
+    config = fetch_auth_config(api, client)
+    _, token_endpoint = _flow_endpoints(config, client, need_device=False)
+    resp = client.post(
+        token_endpoint,
+        data={
+            "grant_type": "refresh_token",
+            "refresh_token": refresh,
+            "client_id": config.client_id,
+        },
+    )
+    if resp.status_code != 200:
+        raise AuthError(f"Session expired for {api}. Run: levain login")
+    payload = resp.json()
+    # Providers may or may not rotate the refresh token; keep the old
+    # one when none comes back, or the next refresh has nothing to use.
+    payload.setdefault("refresh_token", refresh)
+    _store_tokens(api, payload)
+    return str(payload.get("access_token", ""))

levain_recipe/cli.py

@@ -0,0 +1,344 @@
+"""``levain`` — CLI for the Levain recipe toolchain.
+
+Usage:
+    levain validate <file>
+    levain validate-package <dir>
+    levain generate-state <dir>
+    levain mcp
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+import sys
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as pkg_version
+from pathlib import Path
+from typing import Annotated, Any
+
+import typer
+
+from . import LANGUAGE_VERSION
+from ._format import format_errors, format_success
+from ._java import ensure_jar
+from .codegen import generate_state_module
+from .validate import validate, validate_fast, validate_package
+
+app = typer.Typer(
+    name="levain",
+    help="PlusCal spec validation for agent recipes",
+    add_completion=False,
+)
+
+_NO_JAVA_NOTE = (
+    "Java not found — stages 4-5 (TLC model checking) skipped; "
+    "they still run server-side on push."
+)
+
+
+def _version_callback(value: bool) -> None:
+    if value:
+        try:
+            tool = pkg_version("levain-cli")
+        except PackageNotFoundError:
+            tool = "dev"
+        print(f"levain-cli {tool} (recipe language {LANGUAGE_VERSION})")
+        raise typer.Exit()
+
+
+@app.callback()
+def _main(
+    version: Annotated[
+        bool,
+        typer.Option(
+            "--version",
+            callback=_version_callback,
+            is_eager=True,
+            help="Print tool and recipe-language versions.",
+        ),
+    ] = False,
+) -> None:
+    """Levain recipe toolchain."""
+
+
+def _prepare_tlc() -> bool:
+    """Set up TLC (stages 4-5) if Java is available.
+
+    Returns True when full validation can run; prints the
+    degradation note and returns False otherwise.
+    """
+    if shutil.which("java") is None:
+        print(_NO_JAVA_NOTE, file=sys.stderr)
+        return False
+    ensure_jar()
+    return True
+
+
+# ── Commands ──────────────────────────────────────────
+
+
+@app.command("validate")
+def validate_cmd(
+    file: Annotated[
+        Path,
+        typer.Argument(help="Path to .tla file"),
+    ],
+) -> None:
+    """Validate a PlusCal spec file."""
+    if not file.is_file():
+        print(f"Error: file not found: {file}", file=sys.stderr)
+        raise typer.Exit(2)
+
+    source = file.read_text(encoding="utf-8")
+
+    if _prepare_tlc():
+        result = validate(source)
+    else:
+        result = validate_fast(source)
+
+    if isinstance(result, list):
+        print(format_errors(result))
+        raise typer.Exit(1)
+
+    print(format_success(result))
+
+
+@app.command("validate-package")
+def validate_package_cmd(
+    recipe_dir: Annotated[
+        Path,
+        typer.Argument(help="Path to recipe package directory"),
+    ],
+) -> None:
+    """Validate a full recipe package (graph.tla + Python source)."""
+    if not recipe_dir.is_dir():
+        print(
+            f"Error: not a directory: {recipe_dir}",
+            file=sys.stderr,
+        )
+        raise typer.Exit(2)
+
+    _prepare_tlc()
+    result = validate_package(str(recipe_dir))
+    if isinstance(result, list):
+        print(format_errors(result))
+        raise typer.Exit(1)
+
+    print(format_success(result))
+
+
+@app.command("generate-state")
+def generate_state_cmd(
+    recipe_dir: Annotated[
+        Path,
+        typer.Argument(help="Path to recipe package directory"),
+    ],
+) -> None:
+    """Generate the GraphState TypedDict for a recipe.
+
+    Writes ``_state.py`` into the recipe's Python source
+    directory.  Fails if the recipe does not validate.
+    """
+    if not recipe_dir.is_dir():
+        print(
+            f"Error: not a directory: {recipe_dir}",
+            file=sys.stderr,
+        )
+        raise typer.Exit(2)
+
+    _prepare_tlc()
+    result = validate_package(str(recipe_dir))
+    if isinstance(result, list):
+        print(format_errors(result))
+        raise typer.Exit(1)
+
+    src_dir = recipe_dir / "src" / result.module_name
+    if not src_dir.is_dir():
+        print(
+            f"Error: expected src dir at {src_dir}",
+            file=sys.stderr,
+        )
+        raise typer.Exit(2)
+
+    state_path = src_dir / "_state.py"
+    state_path.write_text(generate_state_module(result), encoding="utf-8")
+    print(f"Wrote {state_path}")
+
+
+@app.command("mcp")
+def mcp_cmd() -> None:
+    """Serve the toolchain over stdio MCP.
+
+    For MCP clients that cannot run shell commands (e.g. Claude
+    Desktop); Claude Code should call the CLI directly.
+    """
+    try:
+        from .mcp import server
+    except ImportError:
+        print(
+            "MCP support needs the extra: pip install 'levain-cli[mcp]'",
+            file=sys.stderr,
+        )
+        raise typer.Exit(2) from None
+    server.run()
+
+
+if __name__ == "__main__":
+    app()
+
+
+# ── Authentication ────────────────────────────────────────
+
+
+def _client() -> Any:
+    import httpx
+
+    return httpx.Client(timeout=30.0, follow_redirects=True)
+
+
+@app.command("login")
+def login_cmd() -> None:
+    """Sign in to Levain in a browser.
+
+    Uses the device flow: open the URL shown, enter the code, and the
+    CLI stores a refresh token so git operations can mint their own
+    short-lived credentials.
+    """
+    from .auth import (
+        AuthError,
+        api_base,
+        poll_device_login,
+        start_device_login,
+    )
+
+    api = api_base()
+    try:
+        with _client() as client:
+            prompt = start_device_login(api, client)
+            print(f"Open: {prompt.verification_uri}")
+            if prompt.user_code:
+                print(f"Code:  {prompt.user_code}")
+            print("Waiting for approval...", flush=True)
+            poll_device_login(api, prompt, client)
+    except AuthError as exc:
+        print(f"Error: {exc}", file=sys.stderr)
+        raise typer.Exit(1) from None
+    print(f"Logged in to {api}.")
+
+
+@app.command("logout")
+def logout_cmd() -> None:
+    """Forget the stored Levain credentials."""
+    from .auth import api_base, logout
+
+    api = api_base()
+    print(
+        f"Logged out of {api}." if logout(api) else f"Not logged in to {api}."
+    )
+
+
+@app.command("whoami")
+def whoami_cmd() -> None:
+    """Show who the stored credentials belong to."""
+    from .auth import AuthError, access_token, api_base
+
+    api = api_base()
+    try:
+        with _client() as client:
+            token = access_token(api, client)
+            resp = client.get(
+                f"{api}/api/v1/auth/me",
+                headers={"Authorization": f"Bearer {token}"},
+            )
+    except AuthError as exc:
+        print(f"Error: {exc}", file=sys.stderr)
+        raise typer.Exit(1) from None
+    if resp.status_code != 200:
+        print(f"Error: {api} rejected the session.", file=sys.stderr)
+        raise typer.Exit(1)
+    print(json.dumps(resp.json(), indent=2))
+
+
+@app.command("git-credential")
+def git_credential_cmd(
+    operation: Annotated[
+        str,
+        typer.Argument(help="get, store, or erase (git's helper protocol)"),
+    ],
+) -> None:
+    r"""Git credential helper for Levain recipe repositories.
+
+    Wire it up once and git handles the rest:
+
+        git clone -c credential.helper='!levain git-credential' \
+            -c credential.useHttpPath=true <remote>
+
+    On ``get`` this mints a credential scoped to the repository git
+    asked for, valid for about an hour, so nothing long-lived is ever
+    written to ``.git/config``. ``store`` and ``erase`` are no-ops —
+    there is nothing to persist, and the credential expires on its own.
+    """
+    request = sys.stdin.read()
+    if operation != "get":
+        return
+
+    fields: dict[str, str] = {}
+    for line in request.splitlines():
+        key, _, value = line.partition("=")
+        if key:
+            fields[key.strip()] = value.strip()
+
+    from .auth import AuthError, access_token, api_base, git_host
+
+    # credential.helper is configured repo-wide (see the docstring
+    # above), so git will invoke us for *any* credential request made
+    # inside that repo — e.g. a submodule remote — not just the one we
+    # were cloned from. Refuse anything that isn't our own git host
+    # before minting a real credential for it.
+    if fields.get("protocol") != "https" or fields.get("host") != git_host():
+        print(
+            "levain git-credential: not a Levain recipe URL",
+            file=sys.stderr,
+        )
+        raise typer.Exit(1)
+
+    # git gives us the repo it wants, so the credential is scoped to
+    # exactly that rather than to the whole workspace.
+    slug = _slug_from_git_path(fields.get("path", ""))
+    if not slug:
+        print(
+            "levain git-credential: not a Levain recipe URL",
+            file=sys.stderr,
+        )
+        raise typer.Exit(1)
+
+    api = api_base()
+    try:
+        with _client() as client:
+            token = access_token(api, client)
+            resp = client.post(
+                f"{api}/api/v1/agents/{slug}/git-credential",
+                headers={"Authorization": f"Bearer {token}"},
+            )
+    except AuthError as exc:
+        print(f"levain git-credential: {exc}", file=sys.stderr)
+        raise typer.Exit(1) from None
+    if resp.status_code != 200:
+        print(
+            f"levain git-credential: could not mint a credential for "
+            f"{slug} (HTTP {resp.status_code})",
+            file=sys.stderr,
+        )
+        raise typer.Exit(1)
+    body = resp.json()
+    print(f"username={body.get('username', 'x')}")
+    print(f"password={body['password']}")
+
+
+def _slug_from_git_path(path: str) -> str:
+    """Agent slug from a recipe repo path (``ws-<hex>/<slug>.git``)."""
+    parts = path.strip("/").split("/")
+    if len(parts) != 2 or not parts[0].startswith("ws-"):
+        return ""
+    return parts[1].removesuffix(".git")