@ai-sdk/policy-opa 0.0.0 → 1.0.0-beta.14

package/dist/index.d.ts

@@ -0,0 +1,362 @@
+import { Tool, Context, ToolSet, ModelMessage, InferToolSetContext } from '@ai-sdk/provider-utils';
+import { ToolApprovalConfiguration, ToolApprovalStatus } from 'ai';
+import { LanguageModelV4CallOptions, LanguageModelV4Middleware } from '@ai-sdk/provider';
+
+/**
+ * A generic client for evaluating a policy decision against some external
+ * engine (OPA, Cedar, OpenFGA, a remote HTTP rule service, etc.).
+ *
+ * Each adapter in this package wraps a backend behind this interface so that
+ * the rest of the package (and user code) can switch engines without changing
+ * call sites.
+ */
+interface PolicyClient {
+    /**
+     * Evaluate the given input against the policy identified by `path` and
+     * return the decision payload the engine emitted. Adapters are responsible
+     * for interpreting that payload. For OPA, `opaPolicy` normalizes it into
+     * the SDK's `ToolApprovalStatus` shape.
+     */
+    evaluate<TInput = unknown, TResult = unknown>(path: string, input: TInput): Promise<TResult>;
+}
+
+/**
+ * Narrowed object form of the SDK's `ToolApprovalStatus`.
+ *
+ * The public `ToolApprovalStatus` from `ai` allows both bare strings
+ * (`'approved'`) and the matching object form (`{ type: 'approved' }`). This
+ * package normalizes everything to the object form internally so callers like
+ * `shadow` can rely on a single discriminant.
+ */
+type PolicyDecision = {
+    type: 'approved';
+    reason?: string;
+} | {
+    type: 'denied';
+    reason?: string;
+} | {
+    type: 'user-approval';
+} | {
+    type: 'not-applicable';
+};
+
+/**
+ * Event emitted by {@link shadow} every time the wrapped policy is evaluated.
+ *
+ * Compare `decision` and `effective` to see what would have happened
+ * differently if the policy were enforcing: any record where they disagree
+ * is a tool call that shadow mode is letting through but enforce mode would
+ * have blocked or escalated.
+ */
+interface PolicyDecisionEvent {
+    /** Identifying info from the tool call being evaluated. */
+    toolCall: {
+        toolName: string;
+        toolCallId: string;
+        input: unknown;
+    };
+    /** The decision the wrapped policy returned, normalized to the object form. */
+    decision: PolicyDecision;
+    /** Whether the SDK will act on `decision` (true) or override to allow (false). */
+    enforced: boolean;
+    /** The decision the SDK actually acts on. Equals `decision` when enforcing,
+     *  `{ type: 'approved' }` in shadow mode. */
+    effective: PolicyDecision;
+    /** ISO 8601 timestamp of the evaluation. */
+    timestamp: string;
+}
+/**
+ * Wrap a `toolApproval` in shadow mode so the policy is evaluated and the
+ * decision is reported via `onDecision`, but the SDK is told the call is
+ * approved regardless of what the policy said.
+ *
+ * Use this when you are rolling out a new policy and want to see what it
+ * *would* deny before letting it actually deny anything in production. Wire
+ * `onDecision` to your logger / metrics pipeline, run for a while, inspect
+ * the events where `decision.type !== 'approved'`, fix the policy, then
+ * flip `enforce: true` to graduate.
+ *
+ * @example
+ * ```ts
+ * import { shadow } from '@ai-sdk/policy-opa';
+ * import { opaPolicy, wasmPolicyClient } from '@ai-sdk/policy-opa';
+ *
+ * const client = await wasmPolicyClient({ wasm });
+ *
+ * const toolApproval = shadow(
+ *   opaPolicy({ client, path: 'agent/call/decision' }),
+ *   {
+ *     enforce: process.env.ENFORCE_POLICY === 'true',
+ *     onDecision: (event) => {
+ *       logger.info('policy.decision', {
+ *         tool: event.toolCall.toolName,
+ *         decision: event.decision.type,
+ *         enforced: event.enforced,
+ *         wouldBlock: event.decision.type === 'denied',
+ *       });
+ *     },
+ *   },
+ * );
+ * ```
+ */
+declare function shadow<TOOLS extends Record<string, Tool>, RUNTIME_CONTEXT extends Context | unknown | never = unknown>(approval: ToolApprovalConfiguration<TOOLS & ToolSet, RUNTIME_CONTEXT>, opts?: {
+    /** When true, the SDK acts on the policy's decision. When false (default),
+     *  the SDK is told every call is approved. */
+    enforce?: boolean;
+    /** Invoked once per evaluated tool call with the captured decision. */
+    onDecision?: (event: PolicyDecisionEvent) => void | Promise<void>;
+}): ToolApprovalConfiguration<TOOLS & ToolSet, RUNTIME_CONTEXT>;
+
+type ApprovalLiteralStatus = Extract<ToolApprovalStatus, string>;
+/**
+ * Result returned by {@link wrapMcpTools}: the original tool set plus a
+ * `toolApproval` that falls back to a configurable default for any tool the
+ * supplied approval does not explicitly handle.
+ */
+interface WrappedMcpTools<TOOLS extends Record<string, Tool>, RUNTIME_CONTEXT extends Context | unknown | never> {
+    tools: TOOLS;
+    toolApproval: ToolApprovalConfiguration<TOOLS, RUNTIME_CONTEXT>;
+}
+/**
+ * Apply a fallback approval policy to a discovered tool set so the resulting
+ * configuration is total: every tool in `tools` is gated either by the
+ * supplied `approval` (when it has an opinion) or by `default` (otherwise).
+ *
+ * The motivating case is MCP. An MCP server hands you whatever tools it
+ * exposes, and the agent's effective permissions become the union of the
+ * full server surface. Without a fallback, any tool you forgot to write a
+ * rule for is silently allowed. This helper closes that gap by forcing every
+ * uncovered tool through `default` (which defaults to `user-approval` so the
+ * human is in the loop for anything you didn't think about).
+ *
+ * Works for any tool set, not just MCP-discovered tools. The name reflects
+ * the primary use case rather than a hard constraint.
+ *
+ * @example
+ * ```ts
+ * import { wrapMcpTools } from '@ai-sdk/policy-opa';
+ * import { opaPolicy, wasmPolicyClient } from '@ai-sdk/policy-opa';
+ *
+ * const discovered = await mcpClient.tools();
+ * const client = await wasmPolicyClient({ wasm });
+ *
+ * const { tools, toolApproval } = wrapMcpTools(
+ *   discovered,
+ *   opaPolicy({ client, path: 'agent/call/decision' }),
+ *   { default: 'user-approval' }, // anything OPA does not match needs a human
+ * );
+ *
+ * await generateText({ model, tools, toolApproval, prompt });
+ * ```
+ */
+declare function wrapMcpTools<TOOLS extends Record<string, Tool>, RUNTIME_CONTEXT extends Context | unknown | never = unknown>(tools: TOOLS, approval: ToolApprovalConfiguration<TOOLS, RUNTIME_CONTEXT>, opts?: {
+    default?: ApprovalLiteralStatus;
+}): WrappedMcpTools<TOOLS, RUNTIME_CONTEXT>;
+
+/**
+ * Construct a {@link PolicyClient} that talks to a running OPA server over
+ * HTTP using `@open-policy-agent/opa`.
+ *
+ * The `@open-policy-agent/opa` package is an optional peer dependency; install
+ * it before using this client:
+ *
+ * ```sh
+ * pnpm add @open-policy-agent/opa
+ * ```
+ *
+ * The `url` typically points at `http://localhost:8181` for a locally running
+ * OPA. `headers` is forwarded for Styra DAS / EOPA authentication.
+ */
+declare function httpPolicyClient(opts: {
+    url: string;
+    headers?: Record<string, string>;
+}): PolicyClient;
+
+/**
+ * Normalize an OPA evaluation result into the package's {@link PolicyDecision}
+ * shape.
+ *
+ * Supports two Rego output conventions:
+ *
+ * - **Recommended (explicit):** `{ "decision": "allow" | "deny" | "requires-approval", "reason": string }`.
+ *   Maps to `approved` / `denied` / `user-approval` respectively.
+ *
+ * - **Legacy (boolean):** `{ "allow": boolean, "reason"?: string }`. `true`
+ *   maps to `approved`, `false` to `denied`.
+ *
+ * Unknown shapes and `undefined` are treated as `not-applicable` so that a
+ * Rego rule that does not match any branch defaults to "no opinion" rather
+ * than blocking.
+ */
+declare function normalizeOpaDecision(result: unknown): PolicyDecision;
+
+/**
+ * Default OPA input shape passed to the capability-scoping rule.
+ * Override with `toInput` if your Rego expects a different schema.
+ */
+interface DefaultOpaCapabilityInput {
+    messages: LanguageModelV4CallOptions['prompt'];
+    providerOptions: LanguageModelV4CallOptions['providerOptions'];
+}
+/**
+ * Construct an experimental {@link LanguageModelV4Middleware} that narrows
+ * the `tools` field on every model call to an allowlist returned by OPA.
+ *
+ * Why use this alongside the per-call `toolApproval` gate? Two reasons:
+ *
+ * 1. **Defense in depth.** If a bug or regression slips through the policy
+ *    used by `toolApproval`, the middleware still prevents the model from
+ *    seeing the disallowed tools in the first place.
+ * 2. **Capability disclosure.** The model does not waste tokens describing
+ *    tools it cannot call, and jailbreak attempts get "I don't have access
+ *    to that tool" rather than "[approval denied]".
+ *
+ * The OPA rule at `path` is expected to return either a `string[]` of allowed
+ * tool names, or an object `{ tools: string[] }`. Function tools whose `name`
+ * matches are kept; provider tools whose dotted `id` or bare `name` matches are
+ * kept; everything else is dropped from `params.tools`.
+ *
+ * On a malformed result (or an OPA error), the middleware **fails closed**:
+ * `params.tools` is set to `undefined`, so the model is told it has no tools
+ * available. Misconfiguration should not silently widen capabilities.
+ *
+ * @example
+ * ```ts
+ * import { wrapLanguageModel } from 'ai';
+ * import { wasmPolicyClient, opaCapabilityMiddleware } from '@ai-sdk/policy-opa';
+ *
+ * const client = await wasmPolicyClient({ wasm });
+ *
+ * const wrappedModel = wrapLanguageModel({
+ *   model: anthropic('claude-sonnet-4-5'),
+ *   middleware: opaCapabilityMiddleware({
+ *     client,
+ *     path: 'agent/tools/allowed',
+ *   }),
+ * });
+ * ```
+ */
+declare function opaCapabilityMiddleware(opts: {
+    client: PolicyClient;
+    path: string;
+    toInput?: (args: {
+        messages: LanguageModelV4CallOptions['prompt'];
+        providerOptions: LanguageModelV4CallOptions['providerOptions'];
+    }) => unknown;
+}): LanguageModelV4Middleware;
+
+/**
+ * The default shape passed to the OPA rule as `input` when no `toInput` is
+ * supplied. Rego rules can read `input.tool.name`, `input.args`, and so on.
+ */
+interface DefaultOpaInput {
+    tool: {
+        name: string;
+    };
+    args: unknown;
+    messages: ReadonlyArray<ModelMessage>;
+    runtimeContext: unknown;
+}
+/**
+ * Construct a {@link ToolApprovalConfiguration} backed by an OPA policy.
+ *
+ * The returned generic approval function evaluates the supplied Rego entry
+ * (`path`) for every tool call and maps the result to the SDK's approval
+ * status via {@link normalizeOpaDecision}. Pass the result directly as
+ * `toolApproval` on `generateText` / `streamText` / `ToolLoopAgent`.
+ *
+ * ```ts
+ * import { wasmPolicyClient } from '@ai-sdk/policy-opa';
+ * import { opaPolicy } from '@ai-sdk/policy-opa';
+ *
+ * const client = await wasmPolicyClient({ wasm });
+ * const toolApproval = opaPolicy({ client, path: 'agent/call/decision' });
+ *
+ * await generateText({ model, tools, toolApproval, prompt });
+ * ```
+ *
+ * @param opts.client    The OPA client (HTTP or WASM).
+ * @param opts.path      The Rego entrypoint that returns the decision object.
+ * @param opts.toInput   Optional transformer to shape the OPA input.
+ */
+declare function opaPolicy<TOOLS extends ToolSet = ToolSet, RUNTIME_CONTEXT extends Context | unknown | never = unknown>(opts: {
+    client: PolicyClient;
+    path: string;
+    toInput?: (args: {
+        toolCall: {
+            toolName: string;
+            toolCallId: string;
+            input: unknown;
+        };
+        tools: TOOLS | undefined;
+        toolsContext: InferToolSetContext<TOOLS>;
+        runtimeContext: RUNTIME_CONTEXT;
+        messages: ModelMessage[];
+    }) => unknown;
+}): ToolApprovalConfiguration<TOOLS, RUNTIME_CONTEXT>;
+/**
+ * Optional variant of {@link opaPolicy} that gracefully degrades when no
+ * policy backend is available.
+ *
+ * Returns `undefined` when `client` is `undefined`. Pass that directly as
+ * `toolApproval` and the SDK falls back to its default allow-all behavior.
+ * When `client` is supplied, behaves exactly like {@link opaPolicy}.
+ *
+ * Use this when the policy file is configured per environment (e.g. loaded
+ * in production, absent in local development):
+ *
+ * ```ts
+ * import { readFile } from 'node:fs/promises';
+ * import { optionalOpaPolicy, wasmPolicyClient } from '@ai-sdk/policy-opa';
+ *
+ * const wasm = process.env.POLICY_WASM_PATH
+ *   ? await readFile(process.env.POLICY_WASM_PATH)
+ *   : undefined;
+ * const client = wasm ? await wasmPolicyClient({ wasm }) : undefined;
+ *
+ * const toolApproval = optionalOpaPolicy({
+ *   client,
+ *   path: 'agent/call/decision',
+ * });
+ *
+ * await generateText({ model, tools, toolApproval, prompt });
+ * ```
+ */
+declare function optionalOpaPolicy<TOOLS extends ToolSet = ToolSet, RUNTIME_CONTEXT extends Context | unknown | never = unknown>(opts: {
+    client: PolicyClient | undefined;
+    path: string;
+    toInput?: (args: {
+        toolCall: {
+            toolName: string;
+            toolCallId: string;
+            input: unknown;
+        };
+        tools: TOOLS | undefined;
+        toolsContext: InferToolSetContext<TOOLS>;
+        runtimeContext: RUNTIME_CONTEXT;
+        messages: ModelMessage[];
+    }) => unknown;
+}): ToolApprovalConfiguration<TOOLS, RUNTIME_CONTEXT> | undefined;
+
+/**
+ * Construct a {@link PolicyClient} that evaluates a compiled OPA WASM bundle
+ * in-process using `@open-policy-agent/opa-wasm`.
+ *
+ * The `@open-policy-agent/opa-wasm` package is an optional peer dependency;
+ * install it before using this client:
+ *
+ * ```sh
+ * pnpm add @open-policy-agent/opa-wasm
+ * ```
+ *
+ * The `path` argument to `evaluate(path, input)` is informational. The WASM
+ * bundle is built around a fixed entrypoint at `opa build` time. The path is
+ * recorded for audit logs but does not affect the evaluation.
+ */
+declare function wasmPolicyClient(opts: {
+    wasm: Uint8Array | ArrayBuffer;
+    /** Optional data document bundled into the policy (passed to `setData`). */
+    data?: unknown;
+}): Promise<PolicyClient>;
+
+export { type DefaultOpaCapabilityInput, type DefaultOpaInput, type PolicyClient, type PolicyDecision, type PolicyDecisionEvent, type WrappedMcpTools, httpPolicyClient, normalizeOpaDecision, opaCapabilityMiddleware, opaPolicy, optionalOpaPolicy, shadow, wasmPolicyClient, wrapMcpTools };

package/dist/index.js

@@ -0,0 +1,295 @@
+// src/shadow.ts
+function shadow(approval, opts = {}) {
+  const enforce = opts.enforce === true;
+  const { onDecision } = opts;
+  const wrapped = async (args) => {
+    const raw = await evaluateApproval(approval, args);
+    const decision = normalizePolicyDecision(raw);
+    const effective = enforce ? decision : { type: "approved" };
+    if (onDecision) {
+      const event = {
+        toolCall: {
+          toolName: args.toolCall.toolName,
+          toolCallId: args.toolCall.toolCallId,
+          input: args.toolCall.input
+        },
+        decision,
+        enforced: enforce,
+        effective,
+        timestamp: (/* @__PURE__ */ new Date()).toISOString()
+      };
+      void (async () => {
+        try {
+          await onDecision(event);
+        } catch (e) {
+        }
+      })();
+    }
+    return effective;
+  };
+  return wrapped;
+}
+async function evaluateApproval(approval, args) {
+  if (typeof approval === "function") {
+    return await approval(args);
+  }
+  const map = approval;
+  const perTool = map[args.toolCall.toolName];
+  if (perTool == null) return void 0;
+  if (typeof perTool === "function") {
+    return await perTool(args.toolCall.input, {
+      toolCallId: args.toolCall.toolCallId,
+      messages: args.messages,
+      toolContext: void 0,
+      runtimeContext: args.runtimeContext
+    });
+  }
+  return perTool;
+}
+function normalizePolicyDecision(status) {
+  if (status == null) return { type: "not-applicable" };
+  if (typeof status === "string") {
+    if (status === "approved" || status === "denied" || status === "user-approval" || status === "not-applicable") {
+      return { type: status };
+    }
+    return { type: "not-applicable" };
+  }
+  if (typeof status === "object" && "type" in status) {
+    return status;
+  }
+  return { type: "not-applicable" };
+}
+
+// src/wrap-mcp-tools.ts
+function wrapMcpTools(tools, approval, opts) {
+  var _a, _b;
+  const fallback = (_a = opts == null ? void 0 : opts.default) != null ? _a : "user-approval";
+  if (typeof approval === "function") {
+    const wrapped = async (args) => {
+      const status = await approval(args);
+      return isNotApplicable(status) ? fallback : status;
+    };
+    return { tools, toolApproval: wrapped };
+  }
+  const filled = {};
+  for (const name of Object.keys(tools)) {
+    filled[name] = (_b = approval[name]) != null ? _b : fallback;
+  }
+  return {
+    tools,
+    toolApproval: filled
+  };
+}
+function isNotApplicable(status) {
+  if (status == null) return true;
+  if (status === "not-applicable") return true;
+  if (typeof status === "object" && status !== null) {
+    return status.type === "not-applicable";
+  }
+  return false;
+}
+
+// src/opa/http-policy-client.ts
+function httpPolicyClient(opts) {
+  const { url, headers } = opts;
+  let underlying;
+  async function getUnderlying() {
+    if (underlying) return underlying;
+    let mod;
+    try {
+      mod = await import("@open-policy-agent/opa");
+    } catch (cause) {
+      throw Object.assign(
+        new Error(
+          'Cannot import "@open-policy-agent/opa". Install it as a peer dependency to use httpPolicyClient().'
+        ),
+        { cause }
+      );
+    }
+    underlying = new mod.OPAClient(url, headers ? { headers } : void 0);
+    return underlying;
+  }
+  return {
+    async evaluate(path, input) {
+      const client = await getUnderlying();
+      return await client.evaluate(path, input);
+    }
+  };
+}
+
+// src/opa/normalize-opa-decision.ts
+function normalizeOpaDecision(result) {
+  if (result == null) {
+    return { type: "not-applicable" };
+  }
+  if (typeof result !== "object") {
+    return { type: "not-applicable" };
+  }
+  const record = result;
+  const reason = typeof record.reason === "string" ? record.reason : void 0;
+  if (typeof record.decision === "string") {
+    switch (record.decision) {
+      case "allow":
+        return withReason("approved", reason);
+      case "deny":
+        return withReason("denied", reason);
+      case "requires-approval":
+        return { type: "user-approval" };
+      case "not-applicable":
+        return { type: "not-applicable" };
+    }
+  }
+  if (typeof record.allow === "boolean") {
+    return withReason(record.allow ? "approved" : "denied", reason);
+  }
+  return { type: "not-applicable" };
+}
+function withReason(type, reason) {
+  return reason ? { type, reason } : { type };
+}
+
+// src/opa/evaluate-policy.ts
+async function evaluatePolicy(client, path, input) {
+  try {
+    return { ok: true, result: await client.evaluate(path, input) };
+  } catch (error) {
+    return { ok: false, error };
+  }
+}
+
+// src/opa/opa-capability-middleware.ts
+function opaCapabilityMiddleware(opts) {
+  const { client, path, toInput } = opts;
+  return {
+    specificationVersion: "v4",
+    async transformParams({ params }) {
+      var _a;
+      if (params.tools == null || params.tools.length === 0) {
+        return params;
+      }
+      const input = (_a = toInput == null ? void 0 : toInput({
+        messages: params.prompt,
+        providerOptions: params.providerOptions
+      })) != null ? _a : {
+        messages: params.prompt,
+        providerOptions: params.providerOptions
+      };
+      const outcome = await evaluatePolicy(client, path, input);
+      if (!outcome.ok) {
+        return { ...params, tools: void 0 };
+      }
+      const allowed = extractAllowedNameSet(outcome.result);
+      if (allowed == null) {
+        return { ...params, tools: void 0 };
+      }
+      let removed = false;
+      const filtered = params.tools.filter((t) => {
+        const keep = t.type === "function" ? allowed.has(t.name) : allowed.has(t.id) || allowed.has(t.name);
+        if (!keep) removed = true;
+        return keep;
+      });
+      if (!removed) return params;
+      return { ...params, tools: filtered };
+    }
+  };
+}
+function extractAllowedNameSet(result) {
+  const list = Array.isArray(result) ? result : result != null && typeof result === "object" ? result.tools : void 0;
+  if (!Array.isArray(list)) return null;
+  const out = /* @__PURE__ */ new Set();
+  for (const item of list) {
+    if (typeof item !== "string") return null;
+    out.add(item);
+  }
+  return out;
+}
+
+// src/opa/opa-policy.ts
+function opaPolicy(opts) {
+  const { client, path, toInput } = opts;
+  return async ({
+    toolCall,
+    tools,
+    toolsContext,
+    runtimeContext,
+    messages
+  }) => {
+    var _a;
+    const opaInput = (_a = toInput == null ? void 0 : toInput({ toolCall, tools, toolsContext, runtimeContext, messages })) != null ? _a : {
+      tool: { name: toolCall.toolName },
+      args: toolCall.input,
+      messages,
+      runtimeContext
+    };
+    const outcome = await evaluatePolicy(client, path, opaInput);
+    if (!outcome.ok) {
+      return {
+        type: "denied",
+        reason: `policy evaluation failed: ${errorMessage(outcome.error)}`
+      };
+    }
+    return normalizeOpaDecision(outcome.result);
+  };
+}
+function errorMessage(cause) {
+  if (cause instanceof Error) return cause.message;
+  if (typeof cause === "object" && cause !== null) {
+    try {
+      return JSON.stringify(cause);
+    } catch (e) {
+    }
+  }
+  return String(cause);
+}
+function optionalOpaPolicy(opts) {
+  if (opts.client == null) return void 0;
+  return opaPolicy({
+    client: opts.client,
+    path: opts.path,
+    toInput: opts.toInput
+  });
+}
+
+// src/opa/wasm-policy-client.ts
+async function wasmPolicyClient(opts) {
+  let mod;
+  try {
+    mod = await import("@open-policy-agent/opa-wasm");
+  } catch (cause) {
+    throw Object.assign(
+      new Error(
+        'Cannot import "@open-policy-agent/opa-wasm". Install it as a peer dependency to use wasmPolicyClient().'
+      ),
+      { cause }
+    );
+  }
+  const policy = await mod.loadPolicy(opts.wasm);
+  if (opts.data !== void 0 && typeof policy.setData === "function") {
+    policy.setData(opts.data);
+  }
+  return {
+    async evaluate(_path, input) {
+      const results = policy.evaluate(input);
+      if (!Array.isArray(results) || results.length === 0) {
+        throw Object.assign(
+          new Error(
+            "OPA WASM policy produced no result. Check that the bundle was built with the correct entrypoint (`opa build -t wasm -e <path>`)."
+          ),
+          { input }
+        );
+      }
+      return results[0].result;
+    }
+  };
+}
+export {
+  httpPolicyClient,
+  normalizeOpaDecision,
+  opaCapabilityMiddleware,
+  opaPolicy,
+  optionalOpaPolicy,
+  shadow,
+  wasmPolicyClient,
+  wrapMcpTools
+};
+//# sourceMappingURL=index.js.map