Metadata-Version: 2.4
Name: enclai-decorator
Version: 0.1.0
Summary: @enclai decorator — run any Python function on a remote enclave service via the agent-proxy
Author-email: James Rea <jfgrea27@gmail.com>
Project-URL: Homepage, https://github.com/jfgrea27/enclai
Project-URL: Source, https://github.com/jfgrea27/enclai
Project-URL: Issues, https://github.com/jfgrea27/enclai/issues
Keywords: confidential-computing,enclave,sev-snp,mcp,attestation
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Requires-Python: >=3.11
Description-Content-Type: text/markdown
Provides-Extra: anycode
Requires-Dist: cloudpickle>=3.0; extra == "anycode"
Provides-Extra: test
Requires-Dist: pytest>=8.0; extra == "test"
Requires-Dist: cloudpickle>=3.0; extra == "test"
Provides-Extra: dev
Requires-Dist: ruff>=0.8; extra == "dev"

# enclai — run any Python function on a remote enclave

Decorate a function with `@enclai` and calling it executes on a remote service
instead of locally. The call travels through the **agent-proxy** (which signs
and audits every request) to the **enclai remote service**, which runs it and
returns the result.

```python
from enclai import enclai


@enclai
def foo(a: int, b: int) -> int:
    return a + b


foo(2, 3)  # -> 5, computed on the enclave; the body never runs locally
```

```
caller process            agent-proxy (Rust)              enclai service (Python)
┌───────────────┐  HTTP   ┌──────────────────┐  HTTP      ┌──────────────────────┐
│ @enclai foo() │ ──────► │ /enclai/* backend │ ─────────► │ POST /invoke          │
│  (client)     │ ◄────── │ OpenAPI allowlist │ ◄───────── │ runs foo, returns 5   │
└───────────────┘         │ signs + audits    │            └──────────────────────┘
                          └──────────────────┘
```

## Layout

| File | Role |
|------|------|
| `enclai/client.py` | the `@enclai` decorator + client transport |
| `enclai/server.py` | the remote service that actually runs functions (`enclai-server`) |
| `enclai/registry.py` | registry of decorated functions, for name-based dispatch |
| `enclai/serialization.py` | wire codecs (`json` and `cloudpickle`) |
| `enclai/runner.py` | runs an `async def` function to completion on the sync server |
| `enclai/config.py` | all behaviour is env-var driven |
| `openapi.yaml` | the agent-proxy allowlist for the `/enclai` backend |
| `examples/tasks.py` | example decorated functions |
| `run-demo.sh` | full end-to-end demo through the real agent-proxy |

## The two roles, one module

The *same* module is imported by both sides. Environment variables decide what
the decorator does:

* **client** (default): calling `foo(...)` ships the call to the remote service.
* **server**: set `ENCLAI_ROLE=server`; importing the module registers the
  functions, and the decorator runs them in-process when `/invoke` dispatches.

```bash
# server (the enclave)
ENCLAI_ROLE=server ENCLAI_APP=examples.tasks PORT=9100 enclai-server

# client
ENCLAI_PROXY_URL=http://localhost:8081 ENCLAI_PREFIX=/enclai python my_client.py
```

`./run-demo.sh` wires up the service + agent-proxy + a client and shows the
signed audit entry produced for the call.

## `async def` functions

Decorating an `async def` gives you back a coroutine function, so the call site
is unchanged — you `await` it exactly as you would locally:

```python
@enclai
async def fetch_and_score(url: str) -> float: ...


score = await fetch_and_score("https://…")  # awaited here, executed on the enclave
```

Two things make that work:

* **Client.** The wrapper is itself `async def`. The blocking HTTP request runs
  on a worker thread (`asyncio.to_thread`), so `asyncio.gather(...)` over several
  remote calls really overlaps them instead of serializing on the event loop.
  The caller context (`caller_context(...)`) still rides each call — the thread
  inherits the current context.
* **Server.** `/invoke` runs the coroutine to completion and returns its value,
  in all three modes (registry dispatch, source shipping, cloudpickle). The wire
  protocol is identical for sync and async functions; nothing on the enclave or
  in the audit log distinguishes them.

`ENCLAI_LOCAL=1` / `ENCLAI_ROLE=server` awaits the original function in-process,
as you'd expect. `@enclai(type="test")` tests may be `async def` too — the
`enclai-test` runner awaits them (under plain `pytest`, they need
`pytest-asyncio` or an explicit `asyncio.run`, same as any other async test).

**Async generators are refused** at decoration time: a remote invocation returns
one value, not a stream. Wrap a coroutine that returns the collected results.

## Registering functions on the enclave: `ENCLAI_APP` vs autodiscover

For registry dispatch the server has to *import* your decorated modules so the
`@enclai` decorators run and populate the registry. There are two ways to tell
it what to import — the same functions end up registered either way:

### Manual registration — `ENCLAI_APP`

Name each module (comma/space separated). Explicit, and lets you import a subset.

```bash
ENCLAI_ROLE=server ENCLAI_APP=app.calc,app.more enclai-server
```

### Autodiscover — `ENCLAI_AUTODISCOVER`

Name a **package**; the server walks it with `pkgutil` and imports every module
beneath it, so every decorated function registers with **no list to maintain**.
Point it at your top-level package and forget about it:

```bash
ENCLAI_ROLE=server ENCLAI_AUTODISCOVER=app enclai-server
```

Per-module import failures are logged, not fatal (one broken submodule won't take
the service down), but if autodiscover is requested and *nothing* registers the
server exits non-zero — so a typo'd package name fails the deploy loudly instead
of silently serving zero functions.

### Building the enclave image in CI — `register-enclave.yml`

Autodiscover is what the reusable GitHub Actions workflow
[`.github/workflows/register-enclave.yml`](../../.github/workflows/register-enclave.yml) is built
around. It turns a repo of `@enclai`-decorated Python into a self-contained
enclave image with no hand-maintained registration. It is a two-stage build:

1. **Discover + prune.** In a full build env (your **pinned `requirements.txt`** +
   the whole package) it runs `enclai-discover`, which imports the package and
   keeps only the modules that define `@enclai` functions **plus their
   intra-package import closure** — not the whole package. So the LLM loop, web
   handlers, etc. never enter the attested image; only the registered code does.
2. **Assemble the minimal image.** It installs the same pinned deps and copies
   *only* the pruned source onto the decorator base image (`decorator.yml` builds
   that from `Dockerfile.server`), sets `ENCLAI_AUTODISCOVER`, asserts
   `GET /functions` is non-empty (failing otherwise), then pushes to GHCR.

(This trims *source*, not wheels — deps still come from `requirements.txt`.)

Call it from an app repo (e.g. `auditing-platform-poc`) with a thin caller
workflow:

```yaml
# .github/workflows/enclave.yml  (in the app repo)
name: enclave
on:
  push:
    branches: [main]
  workflow_dispatch:

jobs:
  enclave:
    permissions:
      contents: read
      packages: write
    uses: jfgrea27/enclai/.github/workflows/register-enclave.yml@main
    with:
      package: app                 # top-level package to autodiscover
      requirements: requirements.txt
      service: auditing-platform-poc-enclave
      tests: tests                 # optional: @enclai(type=test) dir, gates the push (see below)
```

The dependency boundary is `requirements.txt` — its pinned versions are what get
installed into the attested image (and, for the cloudpickle codec, must match the
client's Python, which the base image fixes). See the dependency note below.

### Testing enclave functions — `@enclai(type="test")` + `enclai-test`

`@enclai` also marks tests. A parametrized `@enclai(type="test", src="...")`
registers the function as a **test** rather than an invokable one: it is kept out
of `/invoke` and `/functions`, always runs in-process (never dispatched), and is
collected by the `enclai-test` runner. `src` names the function it exercises, for
reporting.

```python
from enclai import enclai


@enclai
def indexed_amount(base, base_index, current_index): ...


@enclai(type="test", src="app.calc.indexed_amount")
def test_indexed_amount():
    assert indexed_amount.local(Decimal("100"), Decimal("100"), Decimal("110")) == Decimal("110")
```

`enclai-test` discovers the same way the server does (`ENCLAI_APP` /
`ENCLAI_AUTODISCOVER`), runs every test, prints a `PASS`/`FAIL` line each, and
exits non-zero if any fail:

```bash
ENCLAI_AUTODISCOVER=app enclai-test
```

The point is **where** it runs: `register-enclave.yml` runs `enclai-test` against
the enclave image, so your functions are validated against the exact dependencies
and Python they run against in production — a failing test fails the image build,
before it is pushed or attested.

Because the pushed image is trimmed to the `@enclai` closure (above), it carries
no test code. So when you point the workflow's `tests:` input at a directory, it
runs them in a **throwaway layer** — `candidate-image + your tests dir` — with
`ENCLAI_AUTODISCOVER="<package> <tests>"`, then discards that layer. The pushed
image stays minimal, but the tests run against an image identical to it. Tests can
therefore live wherever you keep them (e.g. a top-level `tests/` dir); they do not
need to sit inside the enclave package. A decorated test also still runs under
plain `pytest` unchanged — an `@enclai(type="test")` function executes in-process,
so it is just a normal test function to the collector.

## How does it know what code to run? Three modes, increasing power

### 1. Registry dispatch — JSON (default, safe, auditable)

The function code lives **on the enclave**. The service imports your modules
(`ENCLAI_APP=pkg.mod,pkg.other`), the `@enclai` decorators register them, and
the client sends only a **name + JSON args**. Every value is human-readable in
the audit log. Limitation: arguments and the return value must be JSON
(numbers, strings, lists, dicts).

### 2. Source shipping — JSON args, code as text (`ENCLAI_ALLOW_SOURCE=1`)

The client also sends the function's **source text**; the service `exec`s it.
Lets the enclave run a plain `def` it has never seen. Still JSON args. No
closures or captured state (it's just text). Gated because it execs
client-supplied code.

### 3. cloudpickle — "run *any* code" (`ENCLAI_CODEC=cloudpickle`, server `ENCLAI_ALLOW_PICKLE=1`)

The function is serialized **by value** (bytecode, closure cells, defaults,
referenced globals) and args/results are pickled, so this handles what JSON
can't:

* **closures and lambdas** — captured variables travel with the call;
* **classes / arbitrary objects** as arguments *and* return values (a locally
  defined `Vec`, a `numpy` array, a dataclass);
* **methods and `__main__`-defined helpers**.

```bash
pip install "enclai-decorator[anycode]"   # pulls in cloudpickle
# client
ENCLAI_CODEC=cloudpickle ... python my_client.py
# server
ENCLAI_ROLE=server ENCLAI_ALLOW_PICKLE=1 enclai-server
```

Trade-offs: blobs are opaque base64 in the audit log; the enclave does
`pickle.loads` of client bytes (arbitrary code execution — acceptable only
because the enclave is the trust boundary, hence the explicit
`ENCLAI_ALLOW_PICKLE` gate); client and enclave should run the **same Python
version** (pickled bytecode is version-specific).

## What about classes and non-stdlib libraries (numpy, pandas, …)?

This is the important nuance. cloudpickle ships **your code** by value, but it
pickles library code **by reference** — it stores the import path
(`numpy.ndarray`) and re-imports it on the enclave. So:

* **The enclave image is the dependency boundary.** Any third-party library
  your function imports — `numpy`, `pandas`, `torch`, your own published package
  — must be **installed in the enclave container**, at a matching version.
  cloudpickle does *not* ship the library's wheel.
* If the enclave is missing the library, the import fails *there* and surfaces
  on the client as a `RemoteExecutionError` wrapping `ModuleNotFoundError`.
* **Your own pure-Python helper module** can be shipped by value instead of
  installed, with `enclai.ship_module_by_value(mymodule)` (wraps
  `cloudpickle.register_pickle_by_value`). This does **not** work for
  C-extension libraries — those must be installed.

Practical patterns:

* Bake your dependencies into the enclave image (`requirements.txt` / `uv`) —
  the normal, recommended path.
* For arbitrary user code with arbitrary deps, ship a "fat" enclave image with
  the common libraries pre-installed, or accept a per-call dependency manifest
  and install into a sandbox before executing (heavier; not built in here).

## Does it actually run in an enclave?

The enclai service is plain Python and runs wherever you deploy it. In this
project it is meant to run **inside a confidential-computing enclave** (SEV-SNP
/ CoCo), with the **agent-proxy in front of it** providing the trust story:

* the agent-proxy's **audit signing key is released only after SEV-SNP
  attestation** (see `src/agent-proxy/src/main.rs`), so a verifier can prove an
  audited call chain originated from the attested enclave;
* every `/invoke` is **hash-chained and Ed25519-signed** in the audit log;
* the OpenAPI allowlist means only `POST /invoke`, `GET /functions`,
  `GET /healthz` ever reach the service — anything else is rejected (and
  audited) at the proxy.

So the gated `exec`/`pickle.loads` modes are safe *because* the code runs
attested and isolated, with every invocation recorded. The decorator and
service themselves don't perform attestation — they rely on running behind the
attested agent-proxy. Deploying the service to the cherry cluster behind the
already-hardened proxy is the remaining integration step.

## Wire protocol

`POST {proxy}/enclai/invoke`

```jsonc
// registry / source (json codec)
{"function": "examples.tasks.foo", "args": [2, 3], "kwargs": {}}
{"function": "foo", "source": "def foo(a, b):\n    return a + b", "args": [2, 3]}
// cloudpickle codec — values are base64(pickle)
{"function": "foo", "codec": "cloudpickle", "callable": "...", "args": "...", "kwargs": "..."}
```

Responses: `200 {"result": ...}` (or `{"codec":"cloudpickle","result":"<b64>"}`),
`200 {"error": {"type","message","traceback"}}` when the function raised,
`4xx {"error": "..."}` for protocol/allowlist failures.

## Environment variables

| Var | Side | Default | Meaning |
|-----|------|---------|---------|
| `ENCLAI_PROXY_URL` | client | `http://localhost:8081` | agent-proxy base URL |
| `ENCLAI_PREFIX` | client | `/enclai` | proxy backend prefix |
| `ENCLAI_TIMEOUT` | client | `30` | request timeout (s) |
| `ENCLAI_CODEC` | client | `json` | `json` or `cloudpickle` |
| `ENCLAI_SHIP_SOURCE` | client | unset | `1` to send source with json calls |
| `ENCLAI_SHIP_CALLABLE` | client | `1` | `0` = cloudpickle dispatch by name (don't ship the function) |
| `ENCLAI_ROLE` | server | unset | `server` so imports run locally |
| `ENCLAI_APP` | server | unset | modules to import (comma/space sep) |
| `ENCLAI_AUTODISCOVER` | server | unset | package(s) to walk + import (autodiscover) |
| `ENCLAI_ALLOW_SOURCE` | server | unset | `1` to allow source-shipping |
| `ENCLAI_ALLOW_PICKLE` | server | unset | `1` to allow the cloudpickle codec |
| `ENCLAI_REGISTRY_ONLY` | server | unset | `1` = run only the registry; refuse shipped callables |
| `ENCLAI_LOCAL` | both | unset | `1` forces in-process execution |
| `PORT` | server | `9100` | service listen port |

## Tests

```bash
uv run --extra test python -m pytest tests   # registry, codecs, round-trips, deps, autodiscover
./run-demo.sh                   # end-to-end through the real agent-proxy
```
