ai-shield-core 0.1.0 → 0.3.0

package/src/context/wrap-context.ts

@@ -0,0 +1,475 @@
+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);
+}
+
+// ============================================================
+// propagateTrust — Multi-Agent Trust Propagation
+//
+// In a multi-agent pipeline one agent's output becomes the next agent's
+// input. A successful injection in agent A propagates: A summarizes a
+// poisoned document, B reads A's summary and decides, C executes. The
+// 2026 literature calls this multi-agent contagion — and the standard
+// in-context defenses share an attention substrate with the payload, so
+// the only robust handling is to track trust ACROSS the chain and refuse
+// to let a downstream agent treat upstream output as trusted once any
+// link is contaminated.
+//
+// `propagateTrust()` scans one hop (A → B) as `agent-output`, degrades
+// the effective trust tier on any warn/block, and keeps contamination
+// "sticky": pass the returned `hops` back as `priorChain` for the next
+// link so a poisoning at A still marks the C-hop as contaminated even if
+// C's own payload looks clean.
+// ============================================================
+
+export interface AgentHop {
+  /** The agent that PRODUCED the payload entering this hop. */
+  agentId: string;
+  /** Trust tier the payload was treated as at this hop. */
+  trust: TrustTier;
+  /** Scan decision for this hop's payload. */
+  decision: ScanDecision;
+  /** Violations found at this hop. */
+  violations: Violation[];
+}
+
+export interface PropagateTrustOptions {
+  /**
+   * Trust tier of the producing agent's output. Defaults to `untrusted` —
+   * agent output is attacker-influenceable by construction. Only set to
+   * `trusted` for an agent whose output you control end-to-end.
+   */
+  fromTrust?: TrustTier;
+  /**
+   * Chain returned by an earlier `propagateTrust()` call. Pass it to keep
+   * contamination sticky across A→B→C. Omit for the first link.
+   */
+  priorChain?: AgentHop[];
+  /** Ingestion-scanner strictness for the contagion scan. Default `high`. */
+  strictness?: "low" | "medium" | "high";
+}
+
+export interface TrustPropagationResult {
+  /** No contamination anywhere in the chain (including prior hops). */
+  safe: boolean;
+  /** Worst decision across the whole chain — sticky (once block, stays). */
+  decision: ScanDecision;
+  /**
+   * Trust tier the RECEIVING agent should treat the payload as. Degrades to
+   * `untrusted` the moment this hop — or any prior hop — warns or blocks.
+   */
+  effectiveTrust: TrustTier;
+  /** Full chain including this hop. Feed back as `priorChain` for the next. */
+  hops: AgentHop[];
+  /** Every violation across the chain, newest hop last. */
+  violations: Violation[];
+}
+
+/**
+ * Scan one agent-to-agent hand-off and propagate trust along the chain.
+ *
+ * @param payload      The producing agent's output (= consuming agent's input).
+ * @param fromAgentId  Agent that produced `payload`.
+ * @param toAgentId    Agent about to consume `payload`.
+ *
+ * @example
+ * ```ts
+ * import { propagateTrust } from "ai-shield-core";
+ *
+ * // A → B
+ * let chain = await propagateTrust(aOutput, "researcher", "planner");
+ * // B → C, contamination at A stays sticky through to C
+ * chain = await propagateTrust(bOutput, "planner", "executor", {
+ *   priorChain: chain.hops,
+ * });
+ * if (chain.effectiveTrust !== "trusted" && !chain.safe) {
+ *   // an upstream agent was poisoned — do not let the executor act on it
+ *   haltPipeline(chain.violations);
+ * }
+ * ```
+ */
+export async function propagateTrust(
+  payload: string,
+  fromAgentId: string,
+  toAgentId: string,
+  options: PropagateTrustOptions = {},
+): Promise<TrustPropagationResult> {
+  const fromTrust = options.fromTrust ?? "untrusted";
+  const priorChain = options.priorChain ?? [];
+
+  const scanner = new IngestionScanner({
+    strictness: options.strictness ?? "high",
+  });
+  const scanContext: ScanContext = {
+    source: "agent-output",
+    trustTier: fromTrust,
+    agentId: fromAgentId,
+  };
+  const scan = await scanner.scan(payload, scanContext);
+
+  const hopViolations: Violation[] = scan.violations.map((v) => ({
+    ...v,
+    detail: `${v.detail ?? ""} (${fromAgentId}→${toAgentId})`.trim(),
+  }));
+
+  // Was anything upstream already contaminated?
+  const upstreamWorst = priorChain.reduce<ScanDecision>(
+    (worst, h) => (priority(h.decision) > priority(worst) ? h.decision : worst),
+    "allow",
+  );
+  const upstreamContaminated = upstreamWorst !== "allow";
+
+  // This hop's own decision.
+  const hopDecision = scan.decision;
+
+  // Make contamination explicit as a multi-agent violation (distinct from
+  // the per-segment `ingested_injection` the scanner already produced).
+  if (hopDecision !== "allow") {
+    hopViolations.push({
+      type: "trust_propagation",
+      scanner: "trust-chain",
+      score: hopDecision === "block" ? 1.0 : 0.5,
+      threshold: 0.5,
+      message: `Contagion risk in hand-off ${fromAgentId}→${toAgentId}`,
+      detail: `Agent output flagged at this hop`,
+    });
+  } else if (upstreamContaminated) {
+    hopViolations.push({
+      type: "trust_propagation",
+      scanner: "trust-chain",
+      score: 0.5,
+      threshold: 0.5,
+      message: `Payload reaching ${toAgentId} originates from a contaminated chain`,
+      detail: `Upstream contamination is sticky across hops`,
+    });
+  }
+
+  const hop: AgentHop = {
+    agentId: fromAgentId,
+    trust: fromTrust,
+    decision: hopDecision,
+    violations: hopViolations,
+  };
+  const hops = [...priorChain, hop];
+
+  // Worst decision across the full chain (sticky).
+  const chainDecision: ScanDecision =
+    priority(hopDecision) >= priority(upstreamWorst)
+      ? hopDecision
+      : upstreamWorst;
+
+  // Effective trust degrades to untrusted on ANY contamination in the chain.
+  // A clean hand-off from a `trusted` agent with a clean chain stays trusted.
+  const effectiveTrust: TrustTier =
+    chainDecision === "allow" && fromTrust === "trusted"
+      ? "trusted"
+      : "untrusted";
+
+  return {
+    safe: chainDecision === "allow",
+    decision: chainDecision,
+    effectiveTrust,
+    hops,
+    violations: hops.flatMap((h) => h.violations),
+  };
+}

package/src/cost/pricing.ts

@@ -1,8 +1,15 @@
 import type { ModelPricing } from "../types.js";
 
 // ============================================================
-// Model Pricing Table — Updated Feb 2026
-// Prices in USD per 1M tokens
+// Model Pricing Table — Updated June 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).
+//
+// Note: with the Opus 4.7 generation Anthropic dropped the Opus input/output
+// rate from $15/$75 to $5/$25 and serves the 1M context window at standard
+// pricing (no long-context premium). Earlier tables that still list Opus at
+// $15/$75 over-estimate Opus cost by ~3x.
 // ============================================================
 
 export const MODEL_PRICING: Record<string, ModelPricing> = {
@@ -17,16 +24,21 @@ 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 — June 2026 line-up (Fable 5, Opus 4.8/4.7/4.6, Sonnet 4.6, Haiku 4.5)
+  "claude-fable-5": { inputPer1M: 10.0, outputPer1M: 50.0, cachedInputPer1M: 1.0 },
+  "claude-opus-4-8": { inputPer1M: 5.0, outputPer1M: 25.0, cachedInputPer1M: 0.50 },
+  "claude-opus-4-7": { inputPer1M: 5.0, outputPer1M: 25.0, cachedInputPer1M: 0.50 },
+  "claude-opus-4-6": { inputPer1M: 5.0, outputPer1M: 25.0, cachedInputPer1M: 0.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: 1.0, outputPer1M: 5.0, cachedInputPer1M: 0.10 },
 
   // 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 },
+  fable: { inputPer1M: 10.0, outputPer1M: 50.0, cachedInputPer1M: 1.0 },
+  opus: { inputPer1M: 5.0, outputPer1M: 25.0, cachedInputPer1M: 0.50 },
+  sonnet: { inputPer1M: 3.0, outputPer1M: 15.0, cachedInputPer1M: 0.30 },
+  haiku: { inputPer1M: 1.0, outputPer1M: 5.0, cachedInputPer1M: 0.10 },
 };
 
 /** 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);