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

package/src/opa/wasm-policy-client.ts

@@ -0,0 +1,78 @@
+import type { PolicyClient } from '../policy-client';
+
+/**
+ * Loaded OPA WASM bundle. Compiled offline with `opa build -t wasm` and
+ * passed in as bytes.
+ */
+type LoadedPolicy = {
+  evaluate(input: unknown): Array<{ result: unknown }>;
+  /** Some versions expose `setData` for documents bundled at evaluation time. */
+  setData?(data: unknown): void;
+};
+
+/**
+ * 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.
+ */
+export async function wasmPolicyClient(opts: {
+  wasm: Uint8Array | ArrayBuffer;
+  /** Optional data document bundled into the policy (passed to `setData`). */
+  data?: unknown;
+}): Promise<PolicyClient> {
+  type WasmModule = {
+    loadPolicy(wasm: Uint8Array | ArrayBuffer): Promise<LoadedPolicy>;
+  };
+
+  let mod: WasmModule;
+  try {
+    mod =
+      (await import('@open-policy-agent/opa-wasm')) as unknown as WasmModule;
+  } 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 !== undefined && typeof policy.setData === 'function') {
+    policy.setData(opts.data);
+  }
+
+  return {
+    async evaluate(_path, input) {
+      const results = policy.evaluate(input);
+      // The WASM SDK returns an array of `{ result }` entries, one per
+      // top-level expression in the entrypoint. A non-array or empty array
+      // means the entrypoint produced no value (wrong entrypoint at
+      // `opa build` time, an undefined decision rule, or a misbehaving
+      // bundle). Throw rather than return undefined so the caller fails closed
+      // instead of silently treating it as "no opinion".
+      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 the first entry's `result`, matching the HTTP client's
+      // single-decision shape.
+      return results[0].result as never;
+    },
+  };
+}

package/src/policy-client.ts

@@ -0,0 +1,20 @@
+/**
+ * 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.
+ */
+export 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>;
+}

package/src/policy-decision.ts

@@ -0,0 +1,13 @@
+/**
+ * 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.
+ */
+export type PolicyDecision =
+  | { type: 'approved'; reason?: string }
+  | { type: 'denied'; reason?: string }
+  | { type: 'user-approval' }
+  | { type: 'not-applicable' };

package/src/shadow.ts

@@ -0,0 +1,180 @@
+import type {
+  Context,
+  ModelMessage,
+  Tool,
+  ToolSet,
+} from '@ai-sdk/provider-utils';
+import type { ToolApprovalConfiguration } from 'ai';
+import type { PolicyDecision } from './policy-decision';
+
+interface GenericApprovalArgs {
+  toolCall: { toolName: string; toolCallId: string; input: unknown };
+  tools: unknown;
+  toolsContext: unknown;
+  runtimeContext: unknown;
+  messages: ModelMessage[];
+}
+
+/**
+ * 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.
+ */
+export 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',
+ *       });
+ *     },
+ *   },
+ * );
+ * ```
+ */
+export 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> {
+  const enforce = opts.enforce === true;
+  const { onDecision } = opts;
+
+  const wrapped = async (args: GenericApprovalArgs) => {
+    const raw = await evaluateApproval(approval, args);
+    const decision = normalizePolicyDecision(raw);
+
+    const effective: PolicyDecision = enforce ? decision : { type: 'approved' };
+
+    if (onDecision) {
+      const event: PolicyDecisionEvent = {
+        toolCall: {
+          toolName: args.toolCall.toolName,
+          toolCallId: args.toolCall.toolCallId,
+          input: args.toolCall.input,
+        },
+        decision,
+        enforced: enforce,
+        effective,
+        timestamp: new Date().toISOString(),
+      };
+      // Fire-and-forget so a slow or throwing logger does not block the model.
+      void (async () => {
+        try {
+          await onDecision(event);
+        } catch {
+          // Swallow errors from telemetry; enforcement must not depend on it.
+        }
+      })();
+    }
+
+    return effective;
+  };
+
+  // `wrapped` is typed against the erased `GenericApprovalArgs`; TS can't prove
+  // that matches the SDK's per-tool-set function arm, so bridge it here.
+  return wrapped as unknown as ToolApprovalConfiguration<
+    TOOLS & ToolSet,
+    RUNTIME_CONTEXT
+  >;
+}
+
+async function evaluateApproval(
+  approval: unknown,
+  args: GenericApprovalArgs,
+): Promise<unknown> {
+  if (typeof approval === 'function') {
+    return await (
+      approval as (a: GenericApprovalArgs) => Promise<unknown> | unknown
+    )(args);
+  }
+
+  const map = approval as Record<string, unknown>;
+  const perTool = map[args.toolCall.toolName];
+  if (perTool == null) return undefined;
+  if (typeof perTool === 'function') {
+    return await (
+      perTool as (
+        input: unknown,
+        options: {
+          toolCallId: string;
+          messages: ModelMessage[];
+          toolContext: unknown;
+          runtimeContext: unknown;
+        },
+      ) => Promise<unknown> | unknown
+    )(args.toolCall.input, {
+      toolCallId: args.toolCall.toolCallId,
+      messages: args.messages,
+      toolContext: undefined,
+      runtimeContext: args.runtimeContext,
+    });
+  }
+  return perTool;
+}
+
+function normalizePolicyDecision(status: unknown): PolicyDecision {
+  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 as object)) {
+    return status as PolicyDecision;
+  }
+  return { type: 'not-applicable' };
+}

package/src/wrap-mcp-tools.ts

@@ -0,0 +1,97 @@
+import type { Context, Tool } from '@ai-sdk/provider-utils';
+import type { ToolApprovalConfiguration, ToolApprovalStatus } from 'ai';
+
+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.
+ */
+export 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 });
+ * ```
+ */
+export 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> {
+  const fallback: ApprovalLiteralStatus = opts?.default ?? 'user-approval';
+
+  if (typeof approval === 'function') {
+    const wrapped: ToolApprovalConfiguration<
+      TOOLS,
+      RUNTIME_CONTEXT
+    > = async args => {
+      const status = await (
+        approval as (a: typeof args) => Promise<unknown> | unknown
+      )(args);
+      return isNotApplicable(status) ? fallback : (status as never);
+    };
+    return { tools, toolApproval: wrapped };
+  }
+
+  // Per-tool map form: fill in any tool that the user did not list with the
+  // fallback decision. We use `Object.keys(tools)` (not the approval keys) so
+  // the result is total over the discovered surface.
+  const filled = {} as Record<keyof TOOLS, unknown>;
+  for (const name of Object.keys(tools) as Array<keyof TOOLS>) {
+    filled[name] = approval[name] ?? fallback;
+  }
+
+  // `filled` is a plain per-tool map; cast to the SDK's map arm, which TS
+  // can't infer from the erased `unknown` values.
+  return {
+    tools,
+    toolApproval: filled as ToolApprovalConfiguration<TOOLS, RUNTIME_CONTEXT>,
+  };
+}
+
+function isNotApplicable(status: unknown): boolean {
+  if (status == null) return true;
+  if (status === 'not-applicable') return true;
+  if (typeof status === 'object' && status !== null) {
+    return (status as { type?: unknown }).type === 'not-applicable';
+  }
+  return false;
+}