ai-agent-guardrails 0.0.1

package/src/guard-tools.ts

@@ -0,0 +1,158 @@
+import type { GuardToolsOptions, ToolLike, GuardContext } from './types.js';
+import { createDefaultContext, withTimeout } from './utils.js';
+
+/**
+ * Wrap a toolset with guardrails enforcement
+ *
+ * This is the main entry point for the guardrails package. It wraps AI SDK tools
+ * with policy enforcement, budget checks, timeouts, and audit logging.
+ *
+ * @example
+ * ```ts
+ * const tools = guardTools(mcpTools, {
+ *   policy: createSimplePolicy({ requireApprovalForRisk: ['write', 'admin'] }),
+ *   audit: new ConsoleAuditSink(),
+ *   timeoutMs: 10_000,
+ * });
+ * ```
+ */
+export function guardTools<T extends Record<string, ToolLike>>(
+  tools: T,
+  opts: GuardToolsOptions
+): T {
+  const ctx = opts.ctx ?? createDefaultContext();
+  const audit = opts.audit;
+  const timeoutMs = opts.timeoutMs ?? 15_000;
+  const redactor = opts.redactor;
+
+  const wrapped: Record<string, ToolLike> = {};
+
+  for (const [toolName, tool] of Object.entries(tools)) {
+    const originalExecute = tool.execute;
+
+    // Set needsApproval based on policy decision
+    const needsApproval =
+      tool.needsApproval ??
+      (async (input: unknown) => {
+        try {
+          const { risk, reason } = await opts.policy.classify(toolName, input);
+          const decision = await opts.policy.decide({ toolName, input, ctx, risk, reason });
+
+          // Return true if decision requires approval
+          return Boolean((decision as any).needsApproval);
+        } catch (error) {
+          console.error(`[guardTools] Error in needsApproval check for ${toolName}:`, error);
+          // Fail closed: require approval on error
+          return true;
+        }
+      });
+
+    wrapped[toolName] = {
+      ...tool,
+      needsApproval,
+      execute: originalExecute
+        ? async (input: any) => {
+            const timestamp = Date.now();
+
+            // Redact input before logging if redactor is provided
+            const redactedInput = redactor ? redactor.redact(input) : input;
+
+            audit?.emit({
+              type: 'tool_call_attempted',
+              toolName,
+              input: redactedInput,
+              requestId: ctx.requestId,
+              timestamp,
+            });
+
+            // Check tool call budget
+            ctx.toolCalls += 1;
+            if (ctx.toolCalls > ctx.maxToolCalls) {
+              const reason = `Tool budget exceeded (maxToolCalls=${ctx.maxToolCalls})`;
+              audit?.emit({
+                type: 'budget_exceeded',
+                reason,
+                requestId: ctx.requestId,
+                timestamp: Date.now(),
+              });
+              throw new Error(reason);
+            }
+
+            // Check elapsed time budget
+            if (ctx.maxDurationMs) {
+              const elapsed = Date.now() - ctx.startTime;
+              if (elapsed > ctx.maxDurationMs) {
+                const reason = `Time budget exceeded (maxDurationMs=${ctx.maxDurationMs})`;
+                audit?.emit({
+                  type: 'budget_exceeded',
+                  reason,
+                  requestId: ctx.requestId,
+                  timestamp: Date.now(),
+                });
+                throw new Error(reason);
+              }
+            }
+
+            // Classify and decide
+            const { risk, reason } = await opts.policy.classify(toolName, input);
+            const decision = await opts.policy.decide({ toolName, input, ctx, risk, reason });
+
+            // Block if not allowed
+            if (!decision.allow) {
+              audit?.emit({
+                type: 'tool_call_blocked',
+                toolName,
+                reason: decision.reason,
+                requestId: ctx.requestId,
+                timestamp: Date.now(),
+              });
+              throw new Error(`Tool call blocked: ${decision.reason}`);
+            }
+
+            // Log if approval is needed (AI SDK will handle the actual approval flow)
+            if ((decision as any).needsApproval) {
+              audit?.emit({
+                type: 'tool_call_needs_approval',
+                toolName,
+                reason: (decision as any).reason ?? 'approval required',
+                requestId: ctx.requestId,
+                timestamp: Date.now(),
+              });
+            }
+
+            // Execute with timeout
+            const t0 = Date.now();
+            try {
+              const result = await withTimeout(originalExecute(input), timeoutMs);
+              
+              // Redact output if redactor is provided
+              const redactedResult = redactor ? redactor.redact(result) : result;
+              
+              audit?.emit({
+                type: 'tool_call_executed',
+                toolName,
+                durationMs: Date.now() - t0,
+                requestId: ctx.requestId,
+                timestamp: Date.now(),
+              });
+              
+              return result; // Return original result, not redacted (redaction is for logging only)
+            } catch (error: any) {
+              if (error.message?.includes('timed out')) {
+                audit?.emit({
+                  type: 'tool_call_timeout',
+                  toolName,
+                  timeoutMs,
+                  requestId: ctx.requestId,
+                  timestamp: Date.now(),
+                });
+              }
+              throw error;
+            }
+          }
+        : undefined,
+    };
+  }
+
+  return wrapped as T;
+}

package/src/index.ts

@@ -0,0 +1,30 @@
+// Types
+export type {
+  Risk,
+  GuardDecision,
+  GuardContext,
+  GuardPolicy,
+  AuditEvent,
+  AuditSink,
+  GuardToolsOptions,
+  ToolLike,
+  Redactor,
+} from './types.js';
+
+// Core functions
+export { guardTools } from './guard-tools.js';
+export { createDefaultContext, withTimeout } from './utils.js';
+
+// Policy utilities
+export { createSimplePolicy, PolicyBuilder } from './policy.js';
+
+// Audit sinks
+export { InMemoryAuditSink, ConsoleAuditSink, FileAuditSink } from './audit.js';
+
+// Redaction utilities
+export {
+  createRegexRedactor,
+  createDefaultRedactor,
+  createFieldRedactor,
+  composeRedactors,
+} from './redaction.js';

package/src/policy.ts

@@ -0,0 +1,134 @@
+import type { Risk, GuardPolicy, GuardDecision, GuardContext } from './types.js';
+
+/**
+ * Create a simple policy with allowlist/denylist
+ */
+export function createSimplePolicy(options: {
+  allowlist?: string[];
+  denylist?: string[];
+  requireApprovalForRisk?: Risk[];
+}): GuardPolicy {
+  const { allowlist, denylist, requireApprovalForRisk = ['write', 'admin'] } = options;
+
+  return {
+    classify(toolName: string) {
+      // Classify based on tool name patterns
+      const lowered = toolName.toLowerCase();
+
+      if (
+        lowered.includes('delete') ||
+        lowered.includes('remove') ||
+        lowered.includes('destroy') ||
+        lowered.includes('drop')
+      ) {
+        return { risk: 'admin', reason: 'destructive operation' };
+      }
+
+      if (
+        lowered.includes('create') ||
+        lowered.includes('write') ||
+        lowered.includes('update') ||
+        lowered.includes('modify') ||
+        lowered.includes('insert') ||
+        lowered.includes('send') ||
+        lowered.includes('post')
+      ) {
+        return { risk: 'write', reason: 'write operation' };
+      }
+
+      return { risk: 'read', reason: 'read-only operation' };
+    },
+
+    decide({ toolName, risk, ctx }): GuardDecision {
+      // Check denylist first
+      if (denylist && denylist.includes(toolName)) {
+        return { allow: false, reason: `Tool '${toolName}' is in denylist` };
+      }
+
+      // Check allowlist if provided
+      if (allowlist && !allowlist.includes(toolName)) {
+        return { allow: false, reason: `Tool '${toolName}' is not in allowlist` };
+      }
+
+      // Check if approval is required for this risk level
+      if (requireApprovalForRisk.includes(risk)) {
+        return {
+          allow: true,
+          needsApproval: true,
+          reason: `${risk} operation requires approval`,
+        };
+      }
+
+      return { allow: true };
+    },
+  };
+}
+
+/**
+ * Create a composable policy builder
+ */
+export class PolicyBuilder {
+  private classifiers: Array<(toolName: string, input: unknown) => { risk: Risk; reason?: string } | null> = [];
+  private rules: Array<
+    (args: {
+      toolName: string;
+      input: unknown;
+      ctx: GuardContext;
+      risk: Risk;
+      reason?: string;
+    }) => GuardDecision | null
+  > = [];
+
+  /**
+   * Add a classifier function
+   */
+  addClassifier(
+    classifier: (toolName: string, input: unknown) => { risk: Risk; reason?: string } | null
+  ): this {
+    this.classifiers.push(classifier);
+    return this;
+  }
+
+  /**
+   * Add a decision rule
+   */
+  addRule(
+    rule: (args: {
+      toolName: string;
+      input: unknown;
+      ctx: GuardContext;
+      risk: Risk;
+      reason?: string;
+    }) => GuardDecision | null
+  ): this {
+    this.rules.push(rule);
+    return this;
+  }
+
+  /**
+   * Build the final policy
+   */
+  build(): GuardPolicy {
+    return {
+      classify: (toolName: string, input: unknown) => {
+        // Try each classifier in order
+        for (const classifier of this.classifiers) {
+          const result = classifier(toolName, input);
+          if (result) return result;
+        }
+        // Default to read if no classifier matches
+        return { risk: 'read' };
+      },
+
+      decide: args => {
+        // Try each rule in order
+        for (const rule of this.rules) {
+          const decision = rule(args);
+          if (decision) return decision;
+        }
+        // Default to allow if no rule blocks
+        return { allow: true };
+      },
+    };
+  }
+}

package/src/redaction.ts

@@ -0,0 +1,107 @@
+import type { Redactor } from './types.js';
+
+/**
+ * Default patterns for secret detection
+ */
+const DEFAULT_SECRET_PATTERNS = [
+  /\b(sk-[a-zA-Z0-9]{48})\b/g, // OpenAI API keys
+  /\b(sk_live_[a-zA-Z0-9]{24,})\b/g, // Stripe live keys
+  /\b(sk_test_[a-zA-Z0-9]{24,})\b/g, // Stripe test keys
+  /\b([a-zA-Z0-9_-]{40})\b/g, // GitHub tokens (40 chars)
+  /\b(ghp_[a-zA-Z0-9]{36})\b/g, // GitHub personal access tokens
+  /\b(gho_[a-zA-Z0-9]{36})\b/g, // GitHub OAuth tokens
+  /\b(AKIA[0-9A-Z]{16})\b/g, // AWS access keys
+  /-----BEGIN\s+(?:RSA\s+)?PRIVATE\s+KEY-----/g, // Private keys
+  /\b([a-zA-Z0-9_-]{32,})\b/g, // Generic long tokens
+];
+
+/**
+ * Default patterns for PII detection
+ */
+const DEFAULT_PII_PATTERNS = [
+  /\b\d{3}-\d{2}-\d{4}\b/g, // SSN
+  /\b\d{16}\b/g, // Credit card numbers
+  /\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b/g, // Email addresses
+];
+
+/**
+ * Create a regex-based redactor
+ */
+export function createRegexRedactor(patterns: RegExp[], replacement = '[REDACTED]'): Redactor {
+  return {
+    redact(value: unknown): unknown {
+      if (typeof value === 'string') {
+        let redacted = value;
+        for (const pattern of patterns) {
+          redacted = redacted.replace(pattern, replacement);
+        }
+        return redacted;
+      }
+
+      if (Array.isArray(value)) {
+        return value.map(item => this.redact(item));
+      }
+
+      if (value && typeof value === 'object') {
+        const result: Record<string, unknown> = {};
+        for (const [key, val] of Object.entries(value)) {
+          result[key] = this.redact(val);
+        }
+        return result;
+      }
+
+      return value;
+    },
+  };
+}
+
+/**
+ * Create a default redactor with common secret and PII patterns
+ */
+export function createDefaultRedactor(): Redactor {
+  return createRegexRedactor([...DEFAULT_SECRET_PATTERNS, ...DEFAULT_PII_PATTERNS]);
+}
+
+/**
+ * Field-based redactor that redacts specific fields
+ */
+export function createFieldRedactor(fieldsToRedact: string[], replacement = '[REDACTED]'): Redactor {
+  const fieldSet = new Set(fieldsToRedact.map(f => f.toLowerCase()));
+
+  return {
+    redact(value: unknown): unknown {
+      if (Array.isArray(value)) {
+        return value.map(item => this.redact(item));
+      }
+
+      if (value && typeof value === 'object') {
+        const result: Record<string, unknown> = {};
+        for (const [key, val] of Object.entries(value)) {
+          if (fieldSet.has(key.toLowerCase())) {
+            result[key] = replacement;
+          } else {
+            result[key] = this.redact(val);
+          }
+        }
+        return result;
+      }
+
+      return value;
+    },
+  };
+}
+
+/**
+ * Composite redactor that chains multiple redactors
+ */
+export function composeRedactors(...redactors: Redactor[]): Redactor {
+  return {
+    redact(value: unknown): unknown {
+      let result = value;
+      for (const redactor of redactors) {
+        result = redactor.redact(result);
+      }
+      return result;
+    },
+  };
+}

package/src/types.ts

@@ -0,0 +1,117 @@
+/**
+ * Risk classification for tools
+ */
+export type Risk = 'read' | 'write' | 'admin';
+
+/**
+ * Decision outcome from policy evaluation
+ */
+export type GuardDecision =
+  | { allow: true }
+  | { allow: false; reason: string }
+  | { allow: true; needsApproval: true; reason: string };
+
+/**
+ * Context passed through guard execution
+ */
+export type GuardContext = {
+  requestId: string;
+  toolCalls: number;
+  maxToolCalls: number;
+  startTime: number;
+  maxDurationMs?: number;
+};
+
+/**
+ * Tool-like interface that matches AI SDK tool structure
+ */
+export type ToolLike = {
+  description?: string;
+  inputSchema?: unknown;
+  parameters?: unknown;
+  needsApproval?: boolean | ((input: any) => boolean | Promise<boolean>);
+  execute?: (input: any) => Promise<any>;
+};
+
+/**
+ * Policy interface for classification and decision making
+ */
+export type GuardPolicy = {
+  /**
+   * Classify a tool call by risk level
+   */
+  classify: (
+    toolName: string,
+    input: unknown
+  ) => Promise<{ risk: Risk; reason?: string }> | { risk: Risk; reason?: string };
+
+  /**
+   * Decide whether to allow, block, or require approval
+   */
+  decide: (args: {
+    toolName: string;
+    input: unknown;
+    ctx: GuardContext;
+    risk: Risk;
+    reason?: string;
+  }) => Promise<GuardDecision> | GuardDecision;
+};
+
+/**
+ * Audit event types
+ */
+export type AuditEvent =
+  | { type: 'tool_call_attempted'; toolName: string; input: unknown; requestId: string; timestamp: number }
+  | { type: 'tool_call_blocked'; toolName: string; reason: string; requestId: string; timestamp: number }
+  | {
+      type: 'tool_call_needs_approval';
+      toolName: string;
+      reason: string;
+      requestId: string;
+      timestamp: number;
+    }
+  | {
+      type: 'tool_call_executed';
+      toolName: string;
+      durationMs: number;
+      requestId: string;
+      timestamp: number;
+    }
+  | {
+      type: 'tool_call_timeout';
+      toolName: string;
+      timeoutMs: number;
+      requestId: string;
+      timestamp: number;
+    }
+  | {
+      type: 'budget_exceeded';
+      reason: string;
+      requestId: string;
+      timestamp: number;
+    };
+
+/**
+ * Audit sink interface for logging events
+ */
+export type AuditSink = {
+  emit: (event: AuditEvent) => void | Promise<void>;
+};
+
+/**
+ * Options for guardTools wrapper
+ */
+export type GuardToolsOptions = {
+  policy: GuardPolicy;
+  ctx?: GuardContext;
+  audit?: AuditSink;
+  timeoutMs?: number;
+  redactor?: Redactor;
+};
+
+/**
+ * Redactor interface for PII/secret removal
+ */
+export type Redactor = {
+  redact: (value: unknown) => unknown;
+};

package/src/utils.ts

@@ -0,0 +1,47 @@
+import type { GuardContext } from './types.js';
+
+/**
+ * Generate a unique request ID
+ */
+function cryptoRandomId(): string {
+  // Use crypto.randomUUID() in Node 20+ and modern browsers
+  if (typeof globalThis.crypto?.randomUUID === 'function') {
+    return globalThis.crypto.randomUUID();
+  }
+  // Fallback for environments without crypto.randomUUID
+  return `req_${Math.random().toString(16).slice(2)}`;
+}
+
+/**
+ * Create a default guard context with budget limits
+ */
+export function createDefaultContext(requestId?: string): GuardContext {
+  return {
+    requestId: requestId ?? cryptoRandomId(),
+    toolCalls: 0,
+    maxToolCalls: 8,
+    startTime: Date.now(),
+    maxDurationMs: 60_000, // 60 seconds default
+  };
+}
+
+/**
+ * Wrap a promise with a timeout
+ */
+export function withTimeout<T>(promise: Promise<T>, ms: number): Promise<T> {
+  return new Promise((resolve, reject) => {
+    const timeoutId = setTimeout(() => {
+      reject(new Error(`Operation timed out after ${ms}ms`));
+    }, ms);
+
+    promise
+      .then(value => {
+        clearTimeout(timeoutId);
+        resolve(value);
+      })
+      .catch(error => {
+        clearTimeout(timeoutId);
+        reject(error);
+      });
+  });
+}

package/tsconfig.json

@@ -0,0 +1,9 @@
+{
+  "extends": "../../tsconfig.base.json",
+  "compilerOptions": {
+    "outDir": "dist",
+    "rootDir": "src",
+    "incremental": false
+  },
+  "include": ["src"]
+}

package/tsup.config.ts

@@ -0,0 +1,10 @@
+import { defineConfig } from 'tsup';
+
+export default defineConfig({
+  entry: ['src/index.ts'],
+  format: ['esm'],
+  dts: true,
+  sourcemap: true,
+  clean: true,
+  treeshake: true,
+});