ai-shield-core 0.1.0

package/src/cache/lru.ts

@@ -0,0 +1,93 @@
+// ============================================================
+// LRU Cache — O(1) scan result caching with TTL
+// Uses Map insertion-order for LRU eviction
+// ============================================================
+
+export interface LRUCacheConfig {
+  /** Maximum number of cached entries (default: 1000) */
+  maxSize?: number;
+  /** Time-to-live in milliseconds (default: 300_000 = 5 minutes) */
+  ttlMs?: number;
+}
+
+interface CacheEntry<V> {
+  value: V;
+  expiresAt: number;
+}
+
+export class ScanLRUCache<V = unknown> {
+  private cache = new Map<string, CacheEntry<V>>();
+  private readonly maxSize: number;
+  private readonly ttlMs: number;
+
+  constructor(config: LRUCacheConfig = {}) {
+    this.maxSize = config.maxSize ?? 1000;
+    this.ttlMs = config.ttlMs ?? 300_000;
+  }
+
+  /** Get a cached value. Returns undefined on miss or expiry. Promotes to MRU on hit. */
+  get(key: string): V | undefined {
+    const entry = this.cache.get(key);
+    if (!entry) return undefined;
+
+    if (Date.now() > entry.expiresAt) {
+      this.cache.delete(key);
+      return undefined;
+    }
+
+    // Promote to most-recently-used (Map maintains insertion order)
+    this.cache.delete(key);
+    this.cache.set(key, entry);
+    return entry.value;
+  }
+
+  /** Store a value. Evicts LRU entry if at capacity. */
+  set(key: string, value: V): void {
+    // Remove existing to update position
+    this.cache.delete(key);
+
+    // Evict oldest (first in Map) if at capacity
+    if (this.cache.size >= this.maxSize) {
+      const oldestKey = this.cache.keys().next().value as string;
+      this.cache.delete(oldestKey);
+    }
+
+    this.cache.set(key, {
+      value,
+      expiresAt: Date.now() + this.ttlMs,
+    });
+  }
+
+  /** Check if key exists and is not expired */
+  has(key: string): boolean {
+    return this.get(key) !== undefined;
+  }
+
+  /** Delete a specific key */
+  delete(key: string): boolean {
+    return this.cache.delete(key);
+  }
+
+  /** Clear all entries */
+  clear(): void {
+    this.cache.clear();
+  }
+
+  /** Current number of entries (may include expired) */
+  get size(): number {
+    return this.cache.size;
+  }
+
+  /** Remove all expired entries. Returns count of removed entries. */
+  prune(): number {
+    const now = Date.now();
+    let removed = 0;
+    for (const [key, entry] of this.cache) {
+      if (now > entry.expiresAt) {
+        this.cache.delete(key);
+        removed++;
+      }
+    }
+    return removed;
+  }
+}

package/src/cost/anomaly.ts

@@ -0,0 +1,57 @@
+// ============================================================
+// Cost Anomaly Detection — Z-Score based
+// Flags unusual spending patterns
+// ============================================================
+
+export interface AnomalyResult {
+  isAnomaly: boolean;
+  zScore: number;
+  currentValue: number;
+  mean: number;
+  stdDev: number;
+}
+
+/** Detect anomalies using z-score (>2.5 standard deviations = anomaly) */
+export function detectAnomaly(
+  currentValue: number,
+  historicalValues: number[],
+  threshold: number = 2.5,
+): AnomalyResult {
+  if (historicalValues.length < 3) {
+    // Not enough data to determine anomaly
+    return {
+      isAnomaly: false,
+      zScore: 0,
+      currentValue,
+      mean: currentValue,
+      stdDev: 0,
+    };
+  }
+
+  const mean =
+    historicalValues.reduce((a, b) => a + b, 0) / historicalValues.length;
+  const variance =
+    historicalValues.reduce((sum, val) => sum + (val - mean) ** 2, 0) /
+    historicalValues.length;
+  const stdDev = Math.sqrt(variance);
+
+  if (stdDev === 0) {
+    return {
+      isAnomaly: currentValue !== mean,
+      zScore: currentValue === mean ? 0 : Infinity,
+      currentValue,
+      mean,
+      stdDev: 0,
+    };
+  }
+
+  const zScore = (currentValue - mean) / stdDev;
+
+  return {
+    isAnomaly: zScore > threshold,
+    zScore,
+    currentValue,
+    mean,
+    stdDev,
+  };
+}

package/src/cost/pricing.ts

@@ -0,0 +1,58 @@
+import type { ModelPricing } from "../types.js";
+
+// ============================================================
+// Model Pricing Table — Updated Feb 2026
+// Prices in USD per 1M tokens
+// ============================================================
+
+export const MODEL_PRICING: Record<string, ModelPricing> = {
+  // OpenAI
+  "gpt-5.2": { inputPer1M: 2.50, outputPer1M: 10.0 },
+  "gpt-5.1": { inputPer1M: 2.50, outputPer1M: 10.0 },
+  "gpt-5": { inputPer1M: 2.50, outputPer1M: 10.0 },
+  "gpt-4.1": { inputPer1M: 2.00, outputPer1M: 8.00 },
+  "gpt-4o": { inputPer1M: 2.50, outputPer1M: 10.0 },
+  "gpt-4o-mini": { inputPer1M: 0.15, outputPer1M: 0.60 },
+  "o3": { inputPer1M: 10.0, outputPer1M: 40.0 },
+  "o3-mini": { inputPer1M: 1.10, outputPer1M: 4.40 },
+  "o4-mini": { inputPer1M: 1.10, outputPer1M: 4.40 },
+
+  // Anthropic
+  "claude-opus-4-6": { inputPer1M: 15.0, outputPer1M: 75.0 },
+  "claude-sonnet-4-6": { inputPer1M: 3.0, outputPer1M: 15.0 },
+  "claude-haiku-4-5": { inputPer1M: 0.80, outputPer1M: 4.0 },
+
+  // Aliases
+  "gpt-5.2-turbo": { inputPer1M: 2.50, outputPer1M: 10.0 },
+  opus: { inputPer1M: 15.0, outputPer1M: 75.0 },
+  sonnet: { inputPer1M: 3.0, outputPer1M: 15.0 },
+  haiku: { inputPer1M: 0.80, outputPer1M: 4.0 },
+};
+
+/** Get pricing for a model, fallback to gpt-4o-mini rates */
+export function getModelPricing(model: string): ModelPricing {
+  // Try exact match
+  const exact = MODEL_PRICING[model];
+  if (exact) return exact;
+
+  // Try prefix match (e.g., "gpt-4o-2024-08-06" → "gpt-4o")
+  for (const [key, pricing] of Object.entries(MODEL_PRICING)) {
+    if (model.startsWith(key)) return pricing;
+  }
+
+  // Fallback
+  return { inputPer1M: 0.15, outputPer1M: 0.60 };
+}
+
+/** Estimate cost for a given number of tokens */
+export function estimateCost(
+  model: string,
+  inputTokens: number,
+  outputTokens: number,
+): number {
+  const pricing = getModelPricing(model);
+  return (
+    (inputTokens / 1_000_000) * pricing.inputPer1M +
+    (outputTokens / 1_000_000) * pricing.outputPer1M
+  );
+}

package/src/cost/tracker.ts

@@ -0,0 +1,182 @@
+import type {
+  BudgetConfig,
+  BudgetCheckResult,
+  BudgetPeriod,
+  CostRecord,
+} from "../types.js";
+import { estimateCost } from "./pricing.js";
+
+// ============================================================
+// Cost Tracker — Token counting + budget enforcement
+// Uses Redis for distributed budget tracking (optional)
+// Falls back to in-memory for standalone use
+// ============================================================
+
+/** Minimal Redis interface (compatible with ioredis) */
+export interface RedisLike {
+  get(key: string): Promise<string | null>;
+  incrbyfloat(key: string, increment: number): Promise<string>;
+  expire(key: string, seconds: number): Promise<number>;
+}
+
+/** In-memory fallback store */
+class MemoryStore implements RedisLike {
+  private data = new Map<string, { value: string; expiresAt?: number }>();
+
+  async get(key: string): Promise<string | null> {
+    const entry = this.data.get(key);
+    if (!entry) return null;
+    if (entry.expiresAt && Date.now() > entry.expiresAt) {
+      this.data.delete(key);
+      return null;
+    }
+    return entry.value;
+  }
+
+  async incrbyfloat(key: string, increment: number): Promise<string> {
+    const current = parseFloat((await this.get(key)) ?? "0");
+    const newValue = (current + increment).toString();
+    const entry = this.data.get(key);
+    this.data.set(key, { value: newValue, expiresAt: entry?.expiresAt });
+    return newValue;
+  }
+
+  async expire(key: string, seconds: number): Promise<number> {
+    const entry = this.data.get(key);
+    if (!entry) return 0;
+    entry.expiresAt = Date.now() + seconds * 1000;
+    return 1;
+  }
+}
+
+export class CostTracker {
+  private store: RedisLike;
+  private budgets: Map<string, BudgetConfig>;
+  private records: CostRecord[] = [];
+
+  constructor(
+    budgets: Record<string, BudgetConfig> = {},
+    redis?: RedisLike,
+  ) {
+    this.store = redis ?? new MemoryStore();
+    this.budgets = new Map(Object.entries(budgets));
+  }
+
+  /** Check if a request is within budget BEFORE sending to LLM */
+  async checkBudget(
+    entityId: string,
+    model: string,
+    estimatedInputTokens: number,
+    estimatedOutputTokens: number = 500,
+  ): Promise<BudgetCheckResult> {
+    const budget = this.budgets.get(entityId);
+    if (!budget) {
+      return { allowed: true, currentSpend: 0, remainingBudget: Infinity };
+    }
+
+    const estimated = estimateCost(model, estimatedInputTokens, estimatedOutputTokens);
+    const key = this.budgetKey(entityId, budget.period);
+    const currentSpend = parseFloat((await this.store.get(key)) ?? "0");
+
+    if (currentSpend + estimated > budget.hardLimit) {
+      return {
+        allowed: false,
+        currentSpend,
+        remainingBudget: Math.max(0, budget.hardLimit - currentSpend),
+        warning: `Hard budget limit reached: $${currentSpend.toFixed(2)} / $${budget.hardLimit}`,
+      };
+    }
+
+    const warning =
+      currentSpend + estimated > budget.softLimit
+        ? `Approaching budget: $${currentSpend.toFixed(2)} / $${budget.hardLimit} (${Math.round((currentSpend / budget.hardLimit) * 100)}%)`
+        : undefined;
+
+    return {
+      allowed: true,
+      currentSpend,
+      remainingBudget: budget.hardLimit - currentSpend,
+      warning,
+    };
+  }
+
+  /** Record actual cost AFTER receiving response */
+  async recordCost(
+    entityId: string,
+    model: string,
+    inputTokens: number,
+    outputTokens: number,
+  ): Promise<CostRecord> {
+    const cost = estimateCost(model, inputTokens, outputTokens);
+    const record: CostRecord = {
+      entityId,
+      model,
+      inputTokens,
+      outputTokens,
+      cost,
+      timestamp: new Date(),
+    };
+
+    // Update budget counter
+    const budget = this.budgets.get(entityId);
+    if (budget) {
+      const key = this.budgetKey(entityId, budget.period);
+      await this.store.incrbyfloat(key, cost);
+      await this.store.expire(key, this.periodSeconds(budget.period) * 2);
+    }
+
+    // Also update any matching broader budgets (global, etc.)
+    const globalBudget = this.budgets.get("global");
+    if (globalBudget && entityId !== "global") {
+      const globalKey = this.budgetKey("global", globalBudget.period);
+      await this.store.incrbyfloat(globalKey, cost);
+      await this.store.expire(globalKey, this.periodSeconds(globalBudget.period) * 2);
+    }
+
+    this.records.push(record);
+    return record;
+  }
+
+  /** Get current spend for an entity */
+  async getCurrentSpend(entityId: string): Promise<number> {
+    const budget = this.budgets.get(entityId);
+    if (!budget) return 0;
+    const key = this.budgetKey(entityId, budget.period);
+    return parseFloat((await this.store.get(key)) ?? "0");
+  }
+
+  /** Get all recorded costs (for export/audit) */
+  getRecords(): CostRecord[] {
+    return [...this.records];
+  }
+
+  private budgetKey(entityId: string, period: BudgetPeriod): string {
+    const now = new Date();
+    let periodKey: string;
+
+    switch (period) {
+      case "hourly":
+        periodKey = `${now.getUTCFullYear()}-${now.getUTCMonth()}-${now.getUTCDate()}-${now.getUTCHours()}`;
+        break;
+      case "daily":
+        periodKey = `${now.getUTCFullYear()}-${now.getUTCMonth()}-${now.getUTCDate()}`;
+        break;
+      case "monthly":
+        periodKey = `${now.getUTCFullYear()}-${now.getUTCMonth()}`;
+        break;
+    }
+
+    return `ai-shield:cost:${entityId}:${periodKey}`;
+  }
+
+  private periodSeconds(period: BudgetPeriod): number {
+    switch (period) {
+      case "hourly":
+        return 3600;
+      case "daily":
+        return 86400;
+      case "monthly":
+        return 86400 * 31;
+    }
+  }
+}

package/src/index.ts

@@ -0,0 +1,91 @@
+// ============================================================
+// ai-shield-core — Public API
+// ============================================================
+
+// Main class
+export { AIShield } from "./shield.js";
+
+// Scanners (for custom chain building)
+export { HeuristicScanner, type HeuristicConfig } from "./scanner/heuristic.js";
+export { PIIScanner } from "./scanner/pii.js";
+export { ScannerChain, type ChainConfig } from "./scanner/chain.js";
+export { injectCanary, checkCanaryLeak } from "./scanner/canary.js";
+
+// Policy
+export { PolicyEngine, type PolicyPreset } from "./policy/engine.js";
+export { ToolPolicyScanner } from "./policy/tools.js";
+
+// Cost
+export { CostTracker, type RedisLike } from "./cost/tracker.js";
+export { detectAnomaly, type AnomalyResult } from "./cost/anomaly.js";
+export { getModelPricing, estimateCost, MODEL_PRICING } from "./cost/pricing.js";
+
+// Audit
+export { AuditLogger, ConsoleAuditStore, MemoryAuditStore } from "./audit/logger.js";
+export type { AuditStore } from "./audit/types.js";
+
+// Cache
+export { ScanLRUCache, type LRUCacheConfig } from "./cache/lru.js";
+
+// Types (re-export everything)
+export type {
+  // Scanner
+  ScanDecision,
+  ScanResult,
+  ScannerResult,
+  Scanner,
+  ScanContext,
+  Violation,
+  ViolationType,
+  // PII
+  PIIType,
+  PIIAction,
+  PIIEntity,
+  PIIConfig,
+  // Tool
+  ToolCall,
+  ToolPermissions,
+  ToolPolicy,
+  ToolManifestPin,
+  // Cost
+  BudgetPeriod,
+  BudgetConfig,
+  CostEstimate,
+  CostRecord,
+  BudgetCheckResult,
+  ModelPricing,
+  // Audit
+  AuditRecord,
+  AuditConfig,
+  // Config
+  ShieldConfig,
+  InjectionConfig,
+  CostConfig,
+  CacheConfig,
+  ToolConfig,
+  PresetName,
+} from "./types.js";
+
+// --- Convenience function ---
+
+import { AIShield } from "./shield.js";
+import type { ShieldConfig, ScanResult, ScanContext } from "./types.js";
+
+/** Quick scan — one line, maximum protection */
+export async function shield(
+  input: string,
+  configOrContext?: ShieldConfig | ScanContext,
+): Promise<ScanResult> {
+  // Detect if second arg is config or context
+  const isConfig = configOrContext && ("injection" in configOrContext || "pii" in configOrContext || "cost" in configOrContext || "preset" in configOrContext && typeof configOrContext.preset === "string" && !("agentId" in configOrContext));
+
+  const config = isConfig ? (configOrContext as ShieldConfig) : {};
+  const context = isConfig ? {} : (configOrContext as ScanContext) ?? {};
+
+  const instance = new AIShield(config);
+  try {
+    return await instance.scan(input, context);
+  } finally {
+    await instance.close();
+  }
+}

package/src/policy/engine.ts

@@ -0,0 +1,163 @@
+import type { ScanDecision, PresetName, PIIAction } from "../types.js";
+
+// ============================================================
+// Policy Engine — Maps presets to scanner configurations
+// 3 presets: public_website, internal_support, ops_agent
+// ============================================================
+
+export interface PolicyPreset {
+  name: PresetName;
+  injection: {
+    threshold: number;
+    action: ScanDecision;
+  };
+  pii: {
+    action: PIIAction;
+    emailAction: PIIAction;
+    phoneAction: PIIAction;
+    creditCardAction: PIIAction;
+    ibanAction: PIIAction;
+  };
+  tools: {
+    dangerousPatterns: string[];
+    maxChainDepth: number;
+  };
+  cost: {
+    defaultDailyBudget: number;
+    warnAtPercent: number;
+  };
+}
+
+const PRESETS: Record<PresetName, PolicyPreset> = {
+  public_website: {
+    name: "public_website",
+    injection: {
+      threshold: 0.25, // Strictest — public-facing
+      action: "block",
+    },
+    pii: {
+      action: "mask",
+      emailAction: "mask",
+      phoneAction: "mask",
+      creditCardAction: "block",
+      ibanAction: "block",
+    },
+    tools: {
+      dangerousPatterns: [
+        "delete_*",
+        "remove_*",
+        "drop_*",
+        "destroy_*",
+        "admin_*",
+        "execute_*",
+        "send_email",
+        "payment_*",
+        "transfer_*",
+        "write_*",
+        "create_*",
+        "update_*",
+      ],
+      maxChainDepth: 3,
+    },
+    cost: {
+      defaultDailyBudget: 10, // USD
+      warnAtPercent: 80,
+    },
+  },
+
+  internal_support: {
+    name: "internal_support",
+    injection: {
+      threshold: 0.35, // Medium — trusted users
+      action: "block",
+    },
+    pii: {
+      action: "mask",
+      emailAction: "mask",
+      phoneAction: "mask",
+      creditCardAction: "mask",
+      ibanAction: "mask",
+    },
+    tools: {
+      dangerousPatterns: [
+        "delete_*",
+        "remove_*",
+        "drop_*",
+        "destroy_*",
+        "admin_*",
+        "payment_*",
+        "transfer_*",
+      ],
+      maxChainDepth: 5,
+    },
+    cost: {
+      defaultDailyBudget: 50,
+      warnAtPercent: 70,
+    },
+  },
+
+  ops_agent: {
+    name: "ops_agent",
+    injection: {
+      threshold: 0.5, // Relaxed — internal agents
+      action: "warn",
+    },
+    pii: {
+      action: "mask",
+      emailAction: "allow",
+      phoneAction: "allow",
+      creditCardAction: "mask",
+      ibanAction: "mask",
+    },
+    tools: {
+      dangerousPatterns: ["drop_*", "destroy_*", "wipe_*", "shutdown_*"],
+      maxChainDepth: 8,
+    },
+    cost: {
+      defaultDailyBudget: 100,
+      warnAtPercent: 60,
+    },
+  },
+};
+
+export class PolicyEngine {
+  private preset: PolicyPreset;
+
+  constructor(presetName: PresetName = "public_website") {
+    this.preset = PRESETS[presetName];
+  }
+
+  getPreset(): PolicyPreset {
+    return this.preset;
+  }
+
+  getInjectionThreshold(): number {
+    return this.preset.injection.threshold;
+  }
+
+  getPIIAction(type?: string): PIIAction {
+    if (!type) return this.preset.pii.action;
+    const key = `${type}Action` as keyof PolicyPreset["pii"];
+    return (this.preset.pii[key] as PIIAction | undefined) ?? this.preset.pii.action;
+  }
+
+  getDangerousToolPatterns(): string[] {
+    return this.preset.tools.dangerousPatterns;
+  }
+
+  getMaxToolChainDepth(): number {
+    return this.preset.tools.maxChainDepth;
+  }
+
+  getDailyBudget(): number {
+    return this.preset.cost.defaultDailyBudget;
+  }
+
+  static getPresetNames(): PresetName[] {
+    return Object.keys(PRESETS) as PresetName[];
+  }
+
+  static getPreset(name: PresetName): PolicyPreset {
+    return PRESETS[name];
+  }
+}