@wafersecurity/workers-ai

Instrument Cloudflare Workers AI (env.AI binding) calls with Wafer guardrails.

Wafer's gateway works by changing your SDK base_url — but env.AI.run(...) is an in-process binding call with no URL to redirect, and the binding belongs to your account. This wrapper instruments it in place:

• Tier-0 guardrails run locally (PII/secret redaction, blocklist, injection heuristic) → zero added network latency; enforced before the model.
• Policy is centralized — guardrail config is fetched from your Wafer project (cached), so changes in the console apply.
• Analytics are unified — every call is logged to Wafer (async via ctx.waitUntil), so binding traffic shows up in your dashboard alongside proxied traffic.
• Fail-open — any Wafer/config error never breaks your AI call.

Install

npm i @wafersecurity/workers-ai

One-line setup (auto-instrument every env.AI.run)

Wrap your default export once — no per-call changes. Set WAFER_PROJECT and WAFER_API_KEY as Worker secrets and existing env.AI.run(...) calls are guarded automatically (via an env proxy):

import { withWafer } from "@wafersecurity/workers-ai";

const handler = {
  async fetch(req, env, ctx) {
    // unchanged — env.AI is already wrapped
    return Response.json(await env.AI.run("@cf/meta/llama-3.3-70b-instruct", { messages }));
  },
};
export default withWafer(handler);

withWafer wraps every handler — fetch, queue, scheduled, tail — so env.AI is guarded in cron jobs and queue consumers too (no handlers dropped):

const handler = {
  async fetch(req, env, ctx) { return Response.json(await env.AI.run(MODEL, { messages })); },
  async queue(batch, env, ctx) {            // e.g. batch LLM generation
    for (const msg of batch.messages) await env.AI.run(MODEL, { prompt: msg.body.prompt });
  },
  async scheduled(controller, env, ctx) {   // cron
    await env.AI.run(MODEL, { prompt: "daily summary" });
  },
};
export default withWafer(handler);

If WAFER_PROJECT/WAFER_API_KEY aren't set, handlers run untouched.

Streaming

For streaming responses ({ stream: true } → a ReadableStream of SSE), chunks pass through immediately (zero added latency) while the full response is accumulated and logged after the stream ends — so you get complete request/response telemetry for streamed calls. Block-action guardrails (secrets/blocklist) still terminate the stream on a hit; mid-stream redaction is not applied (tokens are already sent) — use non-streaming if you need output redaction.

Telemetry

Each call logs model, decision, findings, tokens and a latency breakdown. By default the wrapper follows the project's content setting (Console → Settings → Log request & response content): if it's on, request/response content is captured too (stored post-guardrail — secrets/PII redacted). Override with log: "metadata" (never send content), log: "content" (always), or log: "off" (no logging; enforcement still runs).

LLM-as-judge

If the project has an LLM judge enabled (console → Guardrails → LLM judge), the wrapper runs it through your own env.AI binding against the configured policy — no extra Wafer round-trip. Violations flag or block per your setting.

Explicit usage

import { wrapAI, WaferBlockedError } from "@wafersecurity/workers-ai";

export default {
  async fetch(request, env, ctx) {
    const ai = wrapAI(env.AI, {
      project: "my-app",
      apiKey: env.WAFER_API_KEY,   // create one in the console → Agents → API keys
      ctx,                          // enables async (non-blocking) logging
    });

    try {
      const res = await ai.run("@cf/meta/llama-3.3-70b-instruct", {
        messages: [{ role: "user", content: "Hello from Wafer" }],
      });
      return Response.json(res);
    } catch (e) {
      if (e instanceof WaferBlockedError) return new Response("Blocked", { status: 403 });
      throw e;
    }
  },
};

ai.run(model, inputs, options) is a drop-in for env.AI.run(...). PII/secrets are redacted in messages / prompt / text before the model sees them; blocked inputs throw WaferBlockedError; outputs are checked too.

Options

| Option | Default | Description | |---|---|---| | project | — | Wafer project id (required) | | apiKey | — | WAFER_API_KEY (required) | | baseUrl | https://wafersecurity.ai | Wafer base URL | | ctx | — | Worker ExecutionContext, for non-blocking logging | | enforceOutput | true | Also guard model output | | configTtlMs | 60000 | How long to cache the project policy |