ai-shield-core 0.1.0 → 0.2.0

package/src/context/wrap-context.ts

@@ -0,0 +1,304 @@
+import type {
+  ContextSegment,
+  IngestionSource,
+  TrustTier,
+  WrappedContext,
+  ScanContext,
+  ScanDecision,
+  Violation,
+} from "../types.js";
+import { createHash } from "node:crypto";
+import { IngestionScanner } from "../scanner/ingestion.js";
+
+// ============================================================
+// wrapContext — Trust-Tier Context Streams
+//
+// The deepest finding of the 2026 prompt-injection literature
+// (Parallax, IPI surveys, OWASP LLM01:2025) is that the LLM cannot
+// reliably distinguish *instruction* from *data* once both share the
+// same attention substrate. The only architecturally robust mitigation
+// is privilege separation: tag every segment with its provenance + trust
+// tier, scan untrusted segments aggressively, and let downstream code
+// decide whether instruction-shaped content from a `web`/`rag`/`tool-desc`
+// segment is allowed to influence behaviour.
+//
+// `wrapContext()` is the ergonomic entry point.
+// ============================================================
+
+/**
+ * Input shape for `wrapContext()`. Each named field is conventional;
+ * pass only what applies.
+ */
+export interface WrapContextInput {
+  /** Developer-controlled prompt. Always `trust: "system"`. */
+  system?: string;
+  /** Direct user message(s). `trust: "untrusted"`, `source: "user"`. */
+  user?: string | string[];
+  /** Retrieved documents. `trust: "untrusted"`, `source: "rag"`. */
+  retrieved?: Array<{ content: string; label?: string } | string>;
+  /** MCP / function tool descriptions about to be exposed to the model. */
+  tools?: Array<{ content: string; label?: string } | string>;
+  /** Stored memory facts. `trust: "untrusted"`, `source: "memory"`. */
+  memory?: Array<{ content: string; label?: string } | string>;
+  /** Scraped / fetched web content. */
+  web?: Array<{ content: string; label?: string } | string>;
+  /** Output from another agent (multi-agent pipelines). */
+  agentOutput?: Array<{ content: string; label?: string } | string>;
+  /**
+   * Promote specific named segments to `"trusted"` (e.g. an internal
+   * knowledge base whose contents you control end-to-end).
+   * Match is by `label` substring, case-insensitive.
+   */
+  trustedLabels?: string[];
+}
+
+/**
+ * Build a `WrappedContext` from typed inputs.
+ *
+ * Trust assignment:
+ * - `system` -> system
+ * - `retrieved`/`tools`/`memory`/`web`/`agent-output` -> untrusted
+ * - `user` -> untrusted (a user is not trusted in this threat model — they
+ *   can also inject; the `untrusted` label means "scan aggressively")
+ * - any segment whose `label` matches one of `trustedLabels` -> trusted
+ *
+ * Trust does NOT mean "skip scanning". It only governs how
+ * `assemblePrompt()` and the per-segment policy decide whether to
+ * include the segment in the final assembled prompt.
+ */
+export function wrapContext(input: WrapContextInput): WrappedContext {
+  const segments: ContextSegment[] = [];
+  const trustedLabels = (input.trustedLabels ?? []).map((s) => s.toLowerCase());
+
+  // Critic H1 — substring match would let an attacker-supplied label
+  // like "untrusted-doc-INTERNAL-kb-poisoned" claim trust because it
+  // CONTAINS the trusted prefix. Match exact or path-anchored only.
+  const isTrustedLabel = (label?: string): boolean => {
+    if (!label) return false;
+    const lc = label.toLowerCase();
+    return trustedLabels.some((tl) => lc === tl || lc.startsWith(tl + "/"));
+  };
+
+  const push = (
+    content: string,
+    source: IngestionSource,
+    trust: TrustTier,
+    label?: string,
+  ): void => {
+    if (typeof content !== "string" || content.length === 0) return;
+    segments.push({
+      source,
+      trust,
+      content,
+      label,
+      contentHash: hashContent(content),
+    });
+  };
+
+  // System: always trust=system. The `source` field is unused for
+  // system segments because `trust === "system"` is the authoritative
+  // signal — Analyst A2 round 1 review. We keep `source: "user"` here
+  // only because `ContextSegment.source` is non-optional; any code that
+  // branches on `seg.source` MUST first check `seg.trust !== "system"`.
+  if (input.system) {
+    push(input.system, "user", "system", "system-prompt");
+  }
+
+  // User messages.
+  if (input.user) {
+    const userInputs = Array.isArray(input.user) ? input.user : [input.user];
+    for (const u of userInputs) {
+      push(u, "user", "untrusted", "user");
+    }
+  }
+
+  // Helper for the array-of-{content,label} groups.
+  const pushGroup = (
+    items: Array<{ content: string; label?: string } | string> | undefined,
+    source: IngestionSource,
+  ): void => {
+    if (!items) return;
+    for (const item of items) {
+      const content = typeof item === "string" ? item : item.content;
+      const label = typeof item === "string" ? undefined : item.label;
+      const trust: TrustTier = isTrustedLabel(label) ? "trusted" : "untrusted";
+      push(content, source, trust, label);
+    }
+  };
+
+  pushGroup(input.retrieved, "rag");
+  pushGroup(input.tools, "tool-desc");
+  pushGroup(input.memory, "memory");
+  pushGroup(input.web, "web");
+  pushGroup(input.agentOutput, "agent-output");
+
+  return { segments };
+}
+
+/**
+ * Scan every segment with the source-specific ingestion profile.
+ * Mutates `ctx` in place by attaching `scanResults` + `decision`,
+ * AND returns the same object for chaining.
+ */
+export async function scanWrappedContext(
+  ctx: WrappedContext,
+  options: { strictness?: "low" | "medium" | "high" } = {},
+): Promise<WrappedContext> {
+  const scanner = new IngestionScanner({
+    strictness: options.strictness ?? "high",
+  });
+  const results: NonNullable<WrappedContext["scanResults"]> = [];
+  let worst: ScanDecision = "allow";
+
+  for (let i = 0; i < ctx.segments.length; i += 1) {
+    const seg = ctx.segments[i]!;
+    // System segments skip the scanner — they're developer-authored and
+    // running the heuristic over a real system prompt would flood with
+    // false positives (system prompts ARE instructions, by definition).
+    if (seg.trust === "system") {
+      results.push({ segmentIndex: i, decision: "allow", violations: [] });
+      continue;
+    }
+
+    const scanContext: ScanContext = {
+      source: seg.source,
+      trustTier: seg.trust,
+    };
+    const r = await scanner.scan(seg.content, scanContext);
+    results.push({
+      segmentIndex: i,
+      decision: r.decision,
+      violations: r.violations,
+    });
+    if (priority(r.decision) > priority(worst)) {
+      worst = r.decision;
+    }
+  }
+
+  ctx.scanResults = results;
+  ctx.decision = worst;
+  return ctx;
+}
+
+/**
+ * Assemble a prompt string respecting tier boundaries.
+ *
+ * Order: `system` → `trusted` retrieved/memory/tool-desc → `user`
+ *      → all remaining `untrusted` segments wrapped in fenced markers.
+ *
+ * Why `trusted` before `user`? Putting developer-marked trusted
+ * context above the user message reduces the chance an untrusted user
+ * prompt re-frames the trusted reference material below it.
+ *
+ * Untrusted segments are wrapped in an explicit fence so a downstream
+ * model has a chance to attend to provenance. This is not a guarantee
+ * (no in-band marker is) but it is the single highest-leverage
+ * mitigation we can apply at the toolkit layer per Anthropic +
+ * OpenAI Model Spec guidance.
+ *
+ * Pass `strictMode: true` to OMIT blocked segments entirely. Default
+ * keeps them but fences them with a `<BLOCKED>` marker so an auditor
+ * can see what was tried.
+ */
+export interface AssembleOptions {
+  strictMode?: boolean;
+  /** Custom fence labels. Defaults are sensible. */
+  fences?: {
+    untrusted?: { open: string; close: string };
+    blocked?: { open: string; close: string };
+  };
+}
+
+export function assemblePrompt(
+  ctx: WrappedContext,
+  options: AssembleOptions = {},
+): string {
+  const fences = {
+    untrusted: options.fences?.untrusted ?? {
+      open: "<UNTRUSTED_CONTENT source=",
+      close: "</UNTRUSTED_CONTENT>",
+    },
+    blocked: options.fences?.blocked ?? {
+      open: "<BLOCKED_CONTENT source=",
+      close: "</BLOCKED_CONTENT>",
+    },
+  };
+
+  // Pre-build a segment→index map ONCE. Avoids O(n²) `indexOf` inside the
+  // assembly loop AND removes a TOCTOU on mutable `ctx.segments` (Critic
+  // H2 + Analyst A4 round 1 review).
+  const segmentIndexMap = new Map<ContextSegment, number>();
+  ctx.segments.forEach((s, i) => segmentIndexMap.set(s, i));
+  const segmentResultMap = new Map<number, NonNullable<WrappedContext["scanResults"]>[number]>();
+  for (const r of ctx.scanResults ?? []) {
+    segmentResultMap.set(r.segmentIndex, r);
+  }
+
+  const ordered: ContextSegment[] = [];
+  // 1. system
+  ordered.push(...ctx.segments.filter((s) => s.trust === "system"));
+  // 2. trusted (retrieved/memory/tool-desc the dev marked as trusted)
+  ordered.push(...ctx.segments.filter((s) => s.trust === "trusted"));
+  // 3. user (untrusted, source="user")
+  ordered.push(
+    ...ctx.segments.filter(
+      (s) => s.source === "user" && s.trust === "untrusted",
+    ),
+  );
+  // 4. all remaining untrusted, preserve original order within group.
+  for (const s of ctx.segments) {
+    if (s.trust === "untrusted" && s.source !== "user") {
+      ordered.push(s);
+    }
+  }
+
+  const parts: string[] = [];
+  for (const seg of ordered) {
+    const segIdx = segmentIndexMap.get(seg) ?? -1;
+    const segResult = segIdx >= 0 ? segmentResultMap.get(segIdx) : undefined;
+    const blocked = segResult?.decision === "block";
+
+    if (blocked) {
+      if (options.strictMode) {
+        // Drop entirely.
+        continue;
+      }
+      parts.push(
+        `${fences.blocked.open}"${seg.source}" label="${seg.label ?? ""}">\n${seg.content}\n${fences.blocked.close}`,
+      );
+      continue;
+    }
+
+    if (seg.trust === "system") {
+      parts.push(seg.content);
+    } else if (seg.trust === "trusted") {
+      parts.push(seg.content);
+    } else if (seg.source === "user" && seg.trust === "untrusted") {
+      // User input keeps its natural shape — fencing every user message
+      // creates more noise than signal.
+      parts.push(seg.content);
+    } else {
+      parts.push(
+        `${fences.untrusted.open}"${seg.source}" label="${seg.label ?? ""}">\n${seg.content}\n${fences.untrusted.close}`,
+      );
+    }
+  }
+
+  return parts.join("\n\n");
+}
+
+function hashContent(content: string): string {
+  return createHash("sha256").update(content).digest("hex");
+}
+
+function priority(d: ScanDecision): number {
+  return d === "block" ? 2 : d === "warn" ? 1 : 0;
+}
+
+/**
+ * Convenience aggregator: violations across all scanned segments.
+ */
+export function flattenViolations(ctx: WrappedContext): Violation[] {
+  if (!ctx.scanResults) return [];
+  return ctx.scanResults.flatMap((r) => r.violations);
+}

package/src/cost/pricing.ts

@@ -1,8 +1,10 @@
 import type { ModelPricing } from "../types.js";
 
 // ============================================================
-// Model Pricing Table — Updated Feb 2026
-// Prices in USD per 1M tokens
+// Model Pricing Table — Updated April 2026
+// Prices in USD per 1M tokens.
+// Includes `cachedInputPer1M` for providers that support prompt caching
+// (Anthropic cache reads land at ~10% of standard input rate).
 // ============================================================
 
 export const MODEL_PRICING: Record<string, ModelPricing> = {
@@ -17,16 +19,18 @@ export const MODEL_PRICING: Record<string, ModelPricing> = {
   "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 },
+  // Anthropic — April 2026 line-up (Opus 4.7, Sonnet 4.6, Haiku 4.5)
+  "claude-opus-4-7": { inputPer1M: 15.0, outputPer1M: 75.0, cachedInputPer1M: 1.50 },
+  "claude-opus-4-6": { inputPer1M: 15.0, outputPer1M: 75.0, cachedInputPer1M: 1.50 },
+  "claude-sonnet-4-6": { inputPer1M: 3.0, outputPer1M: 15.0, cachedInputPer1M: 0.30 },
+  "claude-sonnet-4-5": { inputPer1M: 3.0, outputPer1M: 15.0, cachedInputPer1M: 0.30 },
+  "claude-haiku-4-5": { inputPer1M: 0.80, outputPer1M: 4.0, cachedInputPer1M: 0.08 },
 
   // 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 },
+  opus: { inputPer1M: 15.0, outputPer1M: 75.0, cachedInputPer1M: 1.50 },
+  sonnet: { inputPer1M: 3.0, outputPer1M: 15.0, cachedInputPer1M: 0.30 },
+  haiku: { inputPer1M: 0.80, outputPer1M: 4.0, cachedInputPer1M: 0.08 },
 };
 
 /** Get pricing for a model, fallback to gpt-4o-mini rates */

package/src/cost/tracker.ts

@@ -49,17 +49,32 @@ class MemoryStore implements RedisLike {
   }
 }
 
+export interface CostTrackerOptions {
+  /**
+   * Cap on in-memory CostRecord retention (ring-buffer).
+   * Default: 10_000. Set to 0 to disable record retention entirely
+   * (use this in long-running processes that only care about budget
+   * counters, not per-request records).
+   * Override via env: AI_SHIELD_MAX_RECORDS.
+   */
+  maxRecords?: number;
+}
+
 export class CostTracker {
   private store: RedisLike;
   private budgets: Map<string, BudgetConfig>;
   private records: CostRecord[] = [];
+  private maxRecords: number;
 
   constructor(
     budgets: Record<string, BudgetConfig> = {},
     redis?: RedisLike,
+    options: CostTrackerOptions = {},
   ) {
     this.store = redis ?? new MemoryStore();
     this.budgets = new Map(Object.entries(budgets));
+    const envCap = Number(process.env.AI_SHIELD_MAX_RECORDS);
+    this.maxRecords = options.maxRecords ?? (Number.isFinite(envCap) && envCap >= 0 ? envCap : 10_000);
   }
 
   /** Check if a request is within budget BEFORE sending to LLM */
@@ -133,10 +148,29 @@ export class CostTracker {
       await this.store.expire(globalKey, this.periodSeconds(globalBudget.period) * 2);
     }
 
-    this.records.push(record);
+    this.appendRecord(record);
     return record;
   }
 
+  /**
+   * Append a record with ring-buffer semantics to prevent unbounded memory growth.
+   * When maxRecords is 0, records are not retained.
+   */
+  private appendRecord(record: CostRecord): void {
+    if (this.maxRecords === 0) return;
+    this.records.push(record);
+    if (this.records.length > this.maxRecords) {
+      // Drop oldest entries — O(1) amortized using splice(0, overflow)
+      const overflow = this.records.length - this.maxRecords;
+      this.records.splice(0, overflow);
+    }
+  }
+
+  /** Clear all in-memory records (e.g., after export) */
+  clearRecords(): void {
+    this.records.length = 0;
+  }
+
   /** Get current spend for an entity */
   async getCurrentSpend(entityId: string): Promise<number> {
     const budget = this.budgets.get(entityId);

package/src/index.ts

@@ -10,10 +10,42 @@ 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";
+export {
+  IngestionScanner,
+  scanIngested,
+  trustTierForSource,
+  type IngestionScannerConfig,
+  type IngestionScanResult,
+} from "./scanner/ingestion.js";
+
+// Context / Trust-Tier
+export {
+  wrapContext,
+  scanWrappedContext,
+  assemblePrompt,
+  flattenViolations,
+  type WrapContextInput,
+  type AssembleOptions,
+} from "./context/wrap-context.js";
+
+// Memory Canary / Persistence-Poisoning
+export {
+  mintMemoryCanary,
+  verifyMemoryCanary,
+  rotateMemoryCanary,
+  buildSentinelEntry,
+  bulkVerify,
+  type MintMemoryCanaryOptions,
+} from "./canary/memory.js";
 
 // Policy
 export { PolicyEngine, type PolicyPreset } from "./policy/engine.js";
 export { ToolPolicyScanner } from "./policy/tools.js";
+export {
+  CircuitBreakerRegistry,
+  makeBreakerScope,
+  type CircuitBreakerOptions,
+} from "./policy/circuit-breaker.js";
 
 // Cost
 export { CostTracker, type RedisLike } from "./cost/tracker.js";
@@ -37,6 +69,19 @@ export type {
   ScanContext,
   Violation,
   ViolationType,
+  // Ingestion / Trust-Tier (v0.2)
+  IngestionSource,
+  TrustTier,
+  ContextSegment,
+  WrappedContext,
+  // Memory Canary (v0.2)
+  MemoryCanaryEntry,
+  MemoryCanaryVerification,
+  // Circuit Breaker (v0.2)
+  CircuitState,
+  CircuitBreakerConfig,
+  CircuitBreakerDecision,
+  CounterStoreLike,
   // PII
   PIIType,
   PIIAction,
@@ -71,7 +116,15 @@ export type {
 import { AIShield } from "./shield.js";
 import type { ShieldConfig, ScanResult, ScanContext } from "./types.js";
 
-/** Quick scan — one line, maximum protection */
+/**
+ * Quick scan — one line, maximum protection.
+ *
+ * **Performance warning:** This creates a new AIShield instance on every call.
+ * For production use with multiple calls, create a single `new AIShield(config)`
+ * instance and reuse it — this avoids repeated scanner chain setup and teardown.
+ *
+ * Use `createShieldSingleton()` for a cached version that reuses a single instance.
+ */
 export async function shield(
   input: string,
   configOrContext?: ShieldConfig | ScanContext,
@@ -89,3 +142,31 @@ export async function shield(
     await instance.close();
   }
 }
+
+/**
+ * Create a cached shield function that reuses a single AIShield instance.
+ * Much better performance than `shield()` for repeated calls.
+ *
+ * @example
+ * ```ts
+ * const scan = createShieldSingleton({ injection: { strictness: "high" } });
+ * const r1 = await scan("input 1");
+ * const r2 = await scan("input 2");
+ * // Call scan.close() when done (e.g., on process exit)
+ * await scan.close();
+ * ```
+ */
+export function createShieldSingleton(config: ShieldConfig = {}): {
+  (input: string, context?: ScanContext): Promise<ScanResult>;
+  close(): Promise<void>;
+} {
+  const instance = new AIShield(config);
+
+  const scan = (input: string, context?: ScanContext): Promise<ScanResult> => {
+    return instance.scan(input, context);
+  };
+
+  scan.close = (): Promise<void> => instance.close();
+
+  return scan;
+}