manifold-cli 0.0.39__py3-none-any.whl

manifold_cli/__init__.py

@@ -0,0 +1,3 @@
+from manifold_cli.cli import main
+
+__all__ = ["main"]

manifold_cli/__main__.py

@@ -0,0 +1,4 @@
+from manifold_cli.cli import main
+
+if __name__ == "__main__":
+    main()

manifold_cli/bridge.py

@@ -0,0 +1,70 @@
+"""A byte-identical relay between two `ByteStream`s over any transport.
+
+Used to make a policy server on a plain local TCP socket reachable over our
+transport (tcp/iroh) and back, without either end knowing a bridge sits
+in between. No frame decoding: bytes are copied through verbatim.
+"""
+
+import asyncio
+import contextlib
+
+from manifold_cli.transport import ByteStream, Connector, Listener, TransportError
+
+RELAY_CHUNK_SIZE = 64 * 1024
+
+
+async def pump(src: ByteStream, dst: ByteStream) -> None:
+    """Copy `src` to `dst` until EOF, then close both (unblocks the other pump)."""
+    try:
+        while True:
+            chunk = await src.read(RELAY_CHUNK_SIZE)
+            if not chunk:
+                break
+            try:
+                await dst.write(chunk)
+            except TransportError:
+                # dst was already closed by the other pump's EOF teardown; a
+                # failed write there just means this direction is done too.
+                break
+    finally:
+        # First EOF (or error) tears both directions down. A session that ends
+        # with a clean close loses nothing; a hard truncation mid-flight could
+        # drop unread bytes, which is acceptable here.
+        await src.close()
+        await dst.close()
+
+
+async def splice(a: ByteStream, b: ByteStream) -> None:
+    """Pump a->b and b->a concurrently until either direction hits EOF."""
+    async with asyncio.TaskGroup() as group:
+        group.create_task(pump(a, b))
+        group.create_task(pump(b, a))
+
+
+async def run_bridge(listener: Listener, connector: Connector) -> None:
+    """Accept connections on `listener`, dial `connector`, splice each pair.
+
+    Each accepted stream gets its own spliced pair, so several shards can be
+    bridged at once. On cancellation, outstanding splices and the listener are
+    torn down. Returns silently; nothing noisy on normal shutdown.
+    """
+    tasks: set[asyncio.Task[None]] = set()
+    try:
+        while True:
+            stream = await listener.accept()
+            try:
+                peer = await connector.connect()
+            except BaseException:
+                # A failed dial must not leak the accepted stream: close it so
+                # the peer sees EOF instead of a connection frozen mid-air.
+                await stream.close()
+                raise
+            task = asyncio.create_task(splice(stream, peer))
+            tasks.add(task)
+            task.add_done_callback(tasks.discard)
+    finally:
+        for task in tasks:
+            task.cancel()
+        with contextlib.suppress(Exception):
+            await asyncio.gather(*tasks, return_exceptions=True)
+        await listener.close()

manifold_cli/cli.py

@@ -0,0 +1,54 @@
+from importlib.metadata import version
+
+import typer
+
+from manifold_cli.commands import auth as auth_commands
+from manifold_cli.commands import benchmark as benchmark_commands
+from manifold_cli.commands import organization as organization_commands
+from manifold_cli.commands import policy as policy_commands
+from manifold_cli.commands import run as run_commands
+from manifold_cli.commands import runner as runner_commands
+from manifold_cli.commands import token as token_commands
+from manifold_cli.dependencies import Dependencies, build_dependencies
+
+
+def _version_callback(value: bool) -> None:
+    if value:
+        typer.echo(f"manifold-cli {version('manifold-cli')}")
+        raise typer.Exit()
+
+
+def build_app(deps: Dependencies) -> typer.Typer:
+    app = typer.Typer(
+        name="manifold",
+        help="Manifold CLI.",
+        no_args_is_help=True,
+    )
+
+    @app.callback()
+    def _main(
+        version: bool = typer.Option(
+            False,
+            "--version",
+            callback=_version_callback,
+            is_eager=True,
+            help="Show version and exit.",
+        ),
+    ) -> None:
+        pass
+
+    auth_commands.register(app, deps)
+    organization_commands.register(app, deps)
+    policy_commands.register(app, deps)
+    benchmark_commands.register(app, deps)
+    run_commands.register(app, deps)
+    runner_commands.register(app, deps)
+    token_commands.register(app, deps)
+
+    return app
+
+
+def main() -> None:
+    deps = build_dependencies()
+    app = build_app(deps)
+    app()

manifold_cli/clients/auth_flow.py

@@ -0,0 +1,169 @@
+"""Provide httpx authentication flows for Manifold API clients."""
+
+from collections.abc import AsyncGenerator, Generator
+from http import HTTPStatus
+from typing import Protocol
+
+import httpx
+
+from manifold_cli.models.auth import TokenResponse
+from manifold_cli.models.credentials import Credentials
+from manifold_cli.models.errors import (
+    AccessTokenRejectedError,
+    InvalidAccessTokenError,
+    NoRunnerCredentialError,
+    NotLoggedInError,
+    RunnerCredentialRejectedError,
+    SessionExpiredError,
+    WorkOSAuthError,
+)
+from manifold_cli.models.runner_home import RunnerHome
+
+
+class CredentialStore(Protocol):
+    """Persistent home for the stored credentials WorkOSAuth reads and rotates."""
+
+    def load(self) -> Credentials | None: ...
+
+    def save(self, credentials: Credentials) -> None: ...
+
+    def clear(self) -> None: ...
+
+
+class RunnerHomeStore(Protocol):
+    """The runner home holding the minted token RunnerAuth presents."""
+
+    def load(self) -> RunnerHome | None: ...
+
+
+class TokenRefresher(Protocol):
+    """Exchanges a refresh token for a new token pair, or raises WorkOSAuthError."""
+
+    def refresh(self, refresh_token: str) -> TokenResponse: ...
+
+
+class WorkOSAuth(httpx.Auth):
+    def __init__(self, credentials: CredentialStore, refresher: TokenRefresher):
+        self._credentials = credentials
+        self._refresher = refresher
+
+    def sync_auth_flow(
+        self, request: httpx.Request
+    ) -> Generator[httpx.Request, httpx.Response]:
+        creds = self._credentials.load()
+        if creds is None:
+            raise NotLoggedInError
+        request.headers["Authorization"] = f"Bearer {creds.access_token}"
+        response = yield request
+        if response.status_code != HTTPStatus.UNAUTHORIZED:
+            return
+        rotated = self._refresh(creds)
+        request.headers["Authorization"] = f"Bearer {rotated.access_token}"
+        response = yield request
+        if response.status_code == HTTPStatus.UNAUTHORIZED:
+            raise SessionExpiredError
+
+    def _refresh(self, creds: Credentials) -> Credentials:
+        """Rotate and persist the token pair, or fail terminally.
+
+        Clears the stored credentials and raises SessionExpiredError when the
+        refresh itself is rejected: nothing short of that means the session is
+        dead. A rotation that succeeds is persisted even if the retried request
+        still fails; the caller decides what a lingering 401 means.
+        """
+        try:
+            new_tokens = self._refresher.refresh(creds.refresh_token)
+        except WorkOSAuthError as exc:
+            self._credentials.clear()
+            raise SessionExpiredError from exc
+        rotated = creds.model_copy(
+            update={
+                "access_token": new_tokens.access_token,
+                "refresh_token": new_tokens.refresh_token,
+            }
+        )
+        self._credentials.save(rotated)
+        return rotated
+
+
+class RunnerAuth(httpx.Auth):
+    """Presents a minted runner token on every request.
+
+    There is no refresh half: a runner token is long-lived and revoked rather
+    than rotated, so a 401 is terminal rather than the cue to rotate that it is
+    for `WorkOSAuth`.
+    """
+
+    def __init__(self, home: RunnerHomeStore):
+        self._home = home
+
+    def sync_auth_flow(
+        self, request: httpx.Request
+    ) -> Generator[httpx.Request, httpx.Response]:
+        self._authorise(request)
+        response = yield request
+        self._raise_for_rejection(response)
+
+    async def async_auth_flow(
+        self, request: httpx.Request
+    ) -> AsyncGenerator[httpx.Request, httpx.Response]:
+        self._authorise(request)
+        response = yield request
+        self._raise_for_rejection(response)
+
+    def _authorise(self, request: httpx.Request) -> None:
+        home = self._home.load()
+        if home is None:
+            raise NoRunnerCredentialError
+        request.headers["Authorization"] = f"Bearer {home.credential.token}"
+
+    def _raise_for_rejection(self, response: httpx.Response) -> None:
+        if response.status_code == HTTPStatus.UNAUTHORIZED:
+            raise RunnerCredentialRejectedError
+
+
+class InMemoryRunnerTokenAuth(httpx.Auth):
+    """Send a runner token that is kept in memory.
+
+    `manifold policy serve` uses this. Its token is never written to disk, so
+    `RunnerAuth`, which reads the token from a runner home, does not apply.
+    """
+
+    def __init__(self, token: str):
+        self._token = token
+
+    def sync_auth_flow(
+        self, request: httpx.Request
+    ) -> Generator[httpx.Request, httpx.Response]:
+        request.headers["Authorization"] = f"Bearer {self._token}"
+        response = yield request
+        self._raise_for_rejection(response)
+
+    async def async_auth_flow(
+        self, request: httpx.Request
+    ) -> AsyncGenerator[httpx.Request, httpx.Response]:
+        request.headers["Authorization"] = f"Bearer {self._token}"
+        response = yield request
+        self._raise_for_rejection(response)
+
+    def _raise_for_rejection(self, response: httpx.Response) -> None:
+        if response.status_code == HTTPStatus.UNAUTHORIZED:
+            raise RunnerCredentialRejectedError
+
+
+class StaticTokenAuth(httpx.Auth):
+    """Present one access token without loading or refreshing a session."""
+
+    def __init__(self, token: str):
+        # ensure token contains only printable ASCII chars
+        if any(ord(character) < 0x20 or ord(character) > 0x7E for character in token):
+            raise InvalidAccessTokenError
+        self._token = token
+
+    def sync_auth_flow(
+        self, request: httpx.Request
+    ) -> Generator[httpx.Request, httpx.Response]:
+        request.headers["Authorization"] = f"Bearer {self._token}"
+        response = yield request
+        if response.status_code == HTTPStatus.UNAUTHORIZED:
+            raise AccessTokenRejectedError

manifold_cli/clients/credentials.py

@@ -0,0 +1,53 @@
+import logging
+import os
+import tempfile
+from pathlib import Path
+
+from pydantic import ValidationError
+
+from manifold_cli.models.credentials import Credentials
+
+logger = logging.getLogger(__name__)
+
+
+class CredentialsClient:
+    def __init__(self, path: Path):
+        self._path = path
+
+    def load(self) -> Credentials | None:
+        try:
+            raw = self._path.read_text()
+        except FileNotFoundError:
+            return None
+        try:
+            return Credentials.model_validate_json(raw)
+        except ValidationError:
+            logger.warning("credentials file at %s is malformed; ignoring", self._path)
+            return None
+
+    def save(self, credentials: Credentials) -> None:
+        write_private_file(self._path, credentials.model_dump_json())
+
+    def clear(self) -> None:
+        self._path.unlink(missing_ok=True)
+
+
+def write_private_file(path: Path, payload: str) -> None:
+    """Replace `path` with `payload`, readable only by its owner.
+
+    Written to a temporary file in the same directory and renamed, so a reader
+    never sees a half-written credential and a failed write leaves the previous
+    one intact.
+    """
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fd, tmp_name = tempfile.mkstemp(
+        dir=path.parent, prefix=f".{path.name}.", suffix=".tmp"
+    )
+    try:
+        with os.fdopen(fd, "w") as f:
+            f.write(payload)
+        os.chmod(tmp_name, 0o600)
+        os.replace(tmp_name, path)
+    except BaseException:
+        Path(tmp_name).unlink(missing_ok=True)
+        raise

manifold_cli/clients/docker_engine.py

@@ -0,0 +1,43 @@
+"""Read the Docker Engine API over the Docker daemon's Unix socket.
+
+The `docker` command prints Docker image sizes rounded for display, and the engine
+API reports them in bytes.
+"""
+
+from datetime import timedelta
+
+import httpx
+
+DOCKER_SOCKET_PATH = "/var/run/docker.sock"
+# The Docker daemon computes the Docker image disk usage over every Docker
+# image's layers.
+DOCKER_IMAGE_USAGE_TIMEOUT = timedelta(seconds=30)
+
+
+class DockerEngineClient:
+    def __init__(self, socket_path: str):
+        self._client = httpx.Client(
+            base_url="http://docker",
+            transport=httpx.HTTPTransport(uds=socket_path),
+        )
+
+    def close(self) -> None:
+        self._client.close()
+
+    def get_root_dir(self) -> str:
+        """Return the directory where Docker stores Docker images and containers."""
+        response = self._client.get("/info")
+        response.raise_for_status()
+        return response.json()["DockerRootDir"]
+
+    def get_docker_image_usage(self) -> tuple[int, int]:
+        """Return the number of local Docker images and the bytes their layers use."""
+        response = self._client.get(
+            "/v1.52/system/df",
+            params={"type": "image"},
+            timeout=DOCKER_IMAGE_USAGE_TIMEOUT.total_seconds(),
+        )
+        response.raise_for_status()
+        image_usage = response.json()["ImageUsage"]
+        # The Docker daemon leaves out `TotalCount` and `TotalSize` when they are 0.
+        return image_usage.get("TotalCount", 0), image_usage.get("TotalSize", 0)

manifold_cli/clients/image.py

@@ -0,0 +1,130 @@
+"""Read one file out of a container image, using the `docker` command on $PATH.
+
+`benchmark init` reads an image's rubric this way. `docker create` makes the
+image's files available without starting the container. Starting the container
+would run the image's entrypoint, and a benchmark entrypoint needs a GPU and a
+connection to a policy. `docker cp` then copies the file out.
+
+Only an image already on this machine is read. `docker create` downloads an
+image it cannot find, and a benchmark image is tens of gigabytes. A missing
+image raises ImageInspectError instead. The caller reports that as a file it
+could not read, which is different from an image that does not carry the file.
+"""
+
+import contextlib
+import shutil
+import subprocess
+import tarfile
+from io import BytesIO
+
+# Reading files from a local image takes only as long as the disk needs. These
+# limits exist so that a docker daemon that stops responding cannot hold up a
+# registration that has already succeeded.
+INSPECT_TIMEOUT_SECONDS = 30.0
+COPY_TIMEOUT_SECONDS = 120.0
+
+# The message docker prints when `docker cp` asks for a path the container does
+# not have.
+_MISSING_PATH = "could not find the file"
+
+
+class ImageInspectError(Exception):
+    """The image's files could not be read.
+
+    The message says what stopped the read. An image that can be read but does
+    not carry the path returns None instead. The caller reports these two cases
+    differently: one file could not be checked, the other was checked and is
+    not there.
+    """
+
+
+class DockerImageClient:
+    """Reads files out of local container images."""
+
+    def read_file(self, image: str, path: str) -> bytes | None:
+        """Return `path` from `image`, or None when the image does not carry it.
+
+        Raises ImageInspectError when docker is not installed, when the image
+        is not on this machine, or when the copy fails for any reason other
+        than the path being absent.
+        """
+        if shutil.which("docker") is None:
+            raise ImageInspectError("docker is not on $PATH")
+        self._require_local(image)
+        container = self._create(image)
+        try:
+            return self._copy(container, path)
+        finally:
+            self._remove(container)
+
+    def _require_local(self, image: str) -> None:
+        result = self._run(
+            ["docker", "image", "inspect", image], INSPECT_TIMEOUT_SECONDS
+        )
+        if result.returncode != 0:
+            raise ImageInspectError(f"'{image}' is not on this box")
+
+    def _create(self, image: str) -> str:
+        result = self._run(["docker", "create", image], INSPECT_TIMEOUT_SECONDS)
+        if result.returncode != 0:
+            raise ImageInspectError(_said(result) or f"docker could not open '{image}'")
+        container = result.stdout.decode(errors="replace").strip()
+        if not container:
+            raise ImageInspectError(f"docker named no container for '{image}'")
+        return container
+
+    def _copy(self, container: str, path: str) -> bytes | None:
+        result = self._run(
+            ["docker", "cp", f"{container}:{path}", "-"], COPY_TIMEOUT_SECONDS
+        )
+        if result.returncode != 0:
+            said = _said(result)
+            if _MISSING_PATH in said.lower():
+                return None
+            raise ImageInspectError(said or f"docker could not read {path}")
+        return _only_member(result.stdout, path)
+
+    def _remove(self, container: str) -> None:
+        # This runs in a `finally`, after the read has already happened. A
+        # removal that fails or times out must not turn a successful read into
+        # an error, so a failure here is ignored. Nothing else in this module
+        # removes a container that is left behind.
+        with contextlib.suppress(ImageInspectError):
+            self._run(["docker", "rm", "--force", container], INSPECT_TIMEOUT_SECONDS)
+
+    def _run(
+        self, argv: list[str], timeout: float
+    ) -> subprocess.CompletedProcess[bytes]:
+        try:
+            return subprocess.run(argv, capture_output=True, timeout=timeout)
+        except OSError as exc:
+            raise ImageInspectError(f"could not run {argv[0]}: {exc}") from exc
+        except subprocess.TimeoutExpired as exc:
+            raise ImageInspectError(
+                f"`{' '.join(argv[:2])}` did not finish within {timeout:g} s"
+            ) from exc
+
+
+def _said(result: subprocess.CompletedProcess[bytes]) -> str:
+    return result.stderr.decode(errors="replace").strip()
+
+
+def _only_member(archive: bytes, path: str) -> bytes:
+    """Return the first file in the archive that `docker cp` produced.
+
+    An archive holding no file at all raises an error rather than returning
+    None, because None would say the image has no rubric and hide a broken
+    image. Copying a directory would return the first file inside it, and the
+    caller always names a file rather than a directory.
+    """
+    try:
+        with tarfile.open(fileobj=BytesIO(archive)) as tar:
+            for member in tar:
+                if not member.isfile():
+                    continue
+                extracted = tar.extractfile(member)
+                if extracted is not None:
+                    return extracted.read()
+    except tarfile.TarError as exc:
+        raise ImageInspectError(f"docker streamed no readable archive: {exc}") from exc
+    raise ImageInspectError(f"{path} is not a file in the image")