@curatedmcp/tokenshield-core 0.2.0

package/src/proxy/anthropic-passthrough.ts

@@ -0,0 +1,393 @@
+import { IncomingMessage, ServerResponse } from "node:http";
+import { request as httpsRequest } from "node:https";
+import { request as httpRequest } from "node:http";
+import { URL } from "node:url";
+import { createHash, randomUUID } from "node:crypto";
+import type { RequestRecord, ProxyConfig } from "../types.js";
+import { emptyUsage, dollarsFor } from "../pricing.js";
+import { SSEParser } from "./sse.js";
+import { providerForPath } from "../providers/registry.js";
+import { Pipeline } from "../processors/pipeline.js";
+import type { Processor } from "../processors/types.js";
+import { conversationDedup } from "../processors/conversation-dedup.js";
+import { ResponseCache } from "../processors/response-cache.js";
+
+type RecordSink = (record: RequestRecord) => void;
+
+const HOP_BY_HOP = new Set([
+  "connection",
+  "keep-alive",
+  "proxy-authenticate",
+  "proxy-authorization",
+  "te",
+  "trailer",
+  "transfer-encoding",
+  "upgrade",
+  "host",
+]);
+
+function copyHeaders(src: IncomingMessage["headers"]): Record<string, string> {
+  const out: Record<string, string> = {};
+  for (const [name, value] of Object.entries(src)) {
+    if (value === undefined) continue;
+    if (HOP_BY_HOP.has(name.toLowerCase())) continue;
+    out[name] = Array.isArray(value) ? value.join(", ") : value;
+  }
+  return out;
+}
+
+async function readJsonBody(
+  req: IncomingMessage,
+  limitBytes = 64 * 1024 * 1024,
+): Promise<{ raw: Buffer; parsed: unknown }> {
+  const chunks: Buffer[] = [];
+  let total = 0;
+  for await (const chunk of req) {
+    const buf = chunk as Buffer;
+    total += buf.length;
+    if (total > limitBytes) {
+      throw new Error(`Request body exceeds ${limitBytes} bytes`);
+    }
+    chunks.push(buf);
+  }
+  const raw = Buffer.concat(chunks);
+  let parsed: unknown = null;
+  if (raw.length > 0) {
+    try {
+      parsed = JSON.parse(raw.toString("utf8"));
+    } catch {
+      parsed = null;
+    }
+  }
+  return { raw, parsed };
+}
+
+function conversationFingerprint(parsed: unknown): string {
+  if (!parsed || typeof parsed !== "object") return "_";
+  const obj = parsed as Record<string, unknown>;
+  const sys = obj["system"];
+  const messages = Array.isArray(obj["messages"]) ? obj["messages"] : [];
+  const firstUser = messages.find(
+    (m) => typeof m === "object" && m !== null && (m as Record<string, unknown>)["role"] === "user",
+  );
+  return createHash("sha256")
+    .update(JSON.stringify({ sys, firstUser }))
+    .digest("hex")
+    .slice(0, 16);
+}
+
+function bodySize(json: unknown): number {
+  return Buffer.byteLength(JSON.stringify(json ?? null), "utf8");
+}
+
+/**
+ * Singleton pipeline + cache. One process = one set of state.
+ * Future: per-license configuration once cloud-tier gating ships.
+ */
+const PROCESSORS: Processor[] = [conversationDedup];
+const ENABLED = new Set(PROCESSORS.filter((p) => p.enabledByDefault).map((p) => p.id));
+const PIPELINE = new Pipeline({ processors: PROCESSORS, enabled: ENABLED });
+const RESPONSE_CACHE = new ResponseCache();
+
+export function setProcessorEnabled(id: string, enabled: boolean): void {
+  if (enabled) ENABLED.add(id);
+  else ENABLED.delete(id);
+}
+
+export function getProcessorEnabledIds(): string[] {
+  return Array.from(ENABLED);
+}
+
+export function getResponseCacheStats(): { hits: number; misses: number; entries: number; bytes: number } {
+  return RESPONSE_CACHE.stats();
+}
+
+export async function handleAnthropicRequest(
+  req: IncomingMessage,
+  res: ServerResponse,
+  config: ProxyConfig,
+  sink: RecordSink,
+): Promise<void> {
+  const startedAt = Date.now();
+  const requestId = randomUUID();
+
+  let body: { raw: Buffer; parsed: unknown };
+  try {
+    body = await readJsonBody(req);
+  } catch (err) {
+    res.statusCode = 413;
+    res.setHeader("content-type", "text/plain");
+    res.end(`tokenshield: ${(err as Error).message}`);
+    return;
+  }
+
+  const upstream = new URL(req.url ?? "/", config.upstreamBaseUrl);
+  const provider = providerForPath(upstream.pathname);
+  const isHttps = upstream.protocol === "https:";
+  const requester = isHttps ? httpsRequest : httpRequest;
+
+  // ── 1. Determine model + stream flag from raw body (before any rewrite) ──
+  let model = "unknown";
+  let streamed = false;
+  if (provider !== null) {
+    model = provider.extractModel(body.parsed);
+    streamed = provider.isStreaming(body.parsed);
+  }
+
+  // ── 2. Run request-side processors (fail-open) ──────────────────────────
+  let outboundParsed = body.parsed;
+  let outboundBytes = body.raw;
+  const effects: Array<{ name: string; bytesSaved: number; detail?: Record<string, unknown> }> = [];
+  let bytesRaw = body.raw.length;
+  let bytesSent = body.raw.length;
+
+  if (provider !== null && body.parsed !== null) {
+    const conv = provider.toConversation(body.parsed);
+    if (conv !== null) {
+      const ctx = {
+        providerId: provider.id,
+        conversationFingerprint: conversationFingerprint(body.parsed),
+        inboundBytes: body.raw.length,
+      };
+      // sizeOf measures wire-format (after applyConversation), not in-memory shape
+      const wireSize = (c: typeof conv): number =>
+        bodySize(provider.applyConversation(body.parsed, c));
+      const result = PIPELINE.run(conv, ctx, wireSize);
+      for (const e of result.effects) effects.push(e);
+
+      if (result.effects.length > 0) {
+        outboundParsed = provider.applyConversation(body.parsed, result.conversation);
+        const serialized = Buffer.from(JSON.stringify(outboundParsed), "utf8");
+        if (serialized.length < body.raw.length) {
+          outboundBytes = serialized;
+          bytesSent = serialized.length;
+        }
+      }
+    }
+  }
+
+  // ── 3. Response cache: short-circuit if we have a fresh hit ─────────────
+  if (provider !== null && body.parsed !== null) {
+    const hit = RESPONSE_CACHE.lookup(body.parsed);
+    if (hit !== null) {
+      res.statusCode = hit.status;
+      for (const [k, v] of Object.entries(hit.headers)) {
+        if (HOP_BY_HOP.has(k.toLowerCase())) continue;
+        res.setHeader(k, v);
+      }
+      res.setHeader("x-tokenshield-cache", "hit");
+      res.setHeader("x-tokenshield-cache-age-ms", String(hit.cachedAgoMs));
+      res.end(hit.body);
+
+      const dollarsRaw = dollarsFor(hit.model || model, {
+        inputTokens: hit.usage.inputTokens,
+        outputTokens: hit.usage.outputTokens,
+        cacheCreationInputTokens: 0,
+        cacheReadInputTokens: 0,
+      });
+      try {
+        sink({
+          id: requestId,
+          timestamp: startedAt,
+          model: hit.model || model,
+          endpoint: upstream.pathname,
+          streamed: false,
+          durationMs: Date.now() - startedAt,
+          upstreamStatus: hit.status,
+          upstreamError: null,
+          usageRaw: {
+            inputTokens: hit.usage.inputTokens,
+            outputTokens: hit.usage.outputTokens,
+            cacheCreationInputTokens: 0,
+            cacheReadInputTokens: 0,
+          },
+          // Cached: zero new tokens billed
+          usageSent: emptyUsage(),
+          dollarsRaw,
+          dollarsSent: 0,
+          dollarsSaved: dollarsRaw,
+          processorsApplied: ["response-cache:hit", ...effects.map((e) => e.name)],
+        });
+      } catch {
+        /* sink errors must never break the request */
+      }
+      return;
+    }
+  }
+
+  const headers = copyHeaders(req.headers);
+  headers["host"] = upstream.host;
+  headers["content-length"] = String(outboundBytes.length);
+
+  let upstreamStatus = 0;
+  let upstreamError: string | null = null;
+  let usage = emptyUsage();
+  let modelFromResponse: string | null = null;
+  const cachedHeaders: Record<string, string> = {};
+  const responseBodyChunks: Buffer[] = [];
+
+  const finalize = (): void => {
+    const effectiveModel = modelFromResponse ?? model;
+    const dollarsSent = dollarsFor(effectiveModel, usage);
+
+    // Estimate "raw" cost — what the bill would have been without compression.
+    // We use actual sent input tokens + the ratio of bytes saved at the
+    // request layer. This is honest: it's an estimate, marked as such.
+    const totalBytesSavedReq = Math.max(0, bytesRaw - bytesSent);
+    const ratio = bytesSent > 0 ? totalBytesSavedReq / bytesSent : 0;
+    const estimatedInputTokensRaw = Math.round(usage.inputTokens * (1 + ratio));
+    const usageRaw = {
+      inputTokens: estimatedInputTokensRaw,
+      outputTokens: usage.outputTokens,
+      cacheCreationInputTokens: usage.cacheCreationInputTokens,
+      cacheReadInputTokens: usage.cacheReadInputTokens,
+    };
+    const dollarsRaw = dollarsFor(effectiveModel, usageRaw);
+    const dollarsSaved = Math.max(0, dollarsRaw - dollarsSent);
+
+    const record: RequestRecord = {
+      id: requestId,
+      timestamp: startedAt,
+      model: effectiveModel,
+      endpoint: upstream.pathname,
+      streamed,
+      durationMs: Date.now() - startedAt,
+      upstreamStatus,
+      upstreamError,
+      usageRaw,
+      usageSent: usage,
+      dollarsRaw,
+      dollarsSent,
+      dollarsSaved,
+      processorsApplied: effects.map((e) => e.name),
+    };
+    try {
+      sink(record);
+    } catch {
+      /* never */
+    }
+
+    // Cache the JSON response (no-op if not cacheable)
+    if (!streamed && provider !== null && responseBodyChunks.length > 0) {
+      const buf = Buffer.concat(responseBodyChunks);
+      RESPONSE_CACHE.store(body.parsed, {
+        status: upstreamStatus,
+        headers: cachedHeaders,
+        body: buf,
+        usage: { inputTokens: usage.inputTokens, outputTokens: usage.outputTokens },
+        model: effectiveModel,
+      });
+    }
+  };
+
+  await new Promise<void>((resolve) => {
+    const upstreamReq = requester(
+      {
+        method: req.method ?? "POST",
+        hostname: upstream.hostname,
+        port: upstream.port || (isHttps ? 443 : 80),
+        path: upstream.pathname + upstream.search,
+        headers,
+      },
+      (upstreamRes) => {
+        upstreamStatus = upstreamRes.statusCode ?? 0;
+        res.statusCode = upstreamStatus;
+        for (const [name, value] of Object.entries(upstreamRes.headers)) {
+          if (value === undefined) continue;
+          if (HOP_BY_HOP.has(name.toLowerCase())) continue;
+          const strVal = Array.isArray(value) ? value.join(", ") : value;
+          res.setHeader(name, strVal);
+          cachedHeaders[name] = strVal;
+        }
+        if (effects.length > 0) {
+          res.setHeader("x-tokenshield-processors", effects.map((e) => e.name).join(","));
+        }
+
+        const contentType = String(upstreamRes.headers["content-type"] ?? "");
+        const isSse = contentType.includes("text/event-stream");
+
+        if (isSse) {
+          const parser = new SSEParser();
+          const accum = provider !== null ? provider.createStreamAccumulator() : null;
+          upstreamRes.on("data", (chunk: Buffer) => {
+            res.write(chunk);
+            if (accum !== null) {
+              try {
+                for (const ev of parser.push(chunk.toString("utf8"))) {
+                  accum.observe(ev);
+                }
+              } catch { /* accounting must never break the data path */ }
+            }
+          });
+          upstreamRes.on("end", () => {
+            if (accum !== null) {
+              try {
+                for (const ev of parser.flush()) accum.observe(ev);
+                usage = accum.total();
+                modelFromResponse = accum.model();
+              } catch { /* ignore */ }
+            }
+            res.end();
+            finalize();
+            resolve();
+          });
+          upstreamRes.on("error", (err) => {
+            upstreamError = err.message;
+            res.end();
+            finalize();
+            resolve();
+          });
+        } else {
+          upstreamRes.on("data", (chunk: Buffer) => {
+            responseBodyChunks.push(chunk);
+            res.write(chunk);
+          });
+          upstreamRes.on("end", () => {
+            try {
+              const text = Buffer.concat(responseBodyChunks).toString("utf8");
+              if (text.length > 0 && provider !== null) {
+                const parsed = JSON.parse(text);
+                const u = provider.usageFromResponseJson(parsed);
+                usage = u.usage;
+                modelFromResponse = u.model;
+              }
+            } catch { /* non-JSON or parse failure — leave usage at zero */ }
+            res.end();
+            finalize();
+            resolve();
+          });
+          upstreamRes.on("error", (err) => {
+            upstreamError = err.message;
+            res.end();
+            finalize();
+            resolve();
+          });
+        }
+      },
+    );
+
+    upstreamReq.on("error", (err) => {
+      upstreamError = err.message;
+      if (!res.headersSent) {
+        res.statusCode = 502;
+        res.setHeader("content-type", "application/json");
+        res.end(
+          JSON.stringify({
+            type: "error",
+            error: {
+              type: "tokenshield_upstream_error",
+              message: `Failed to reach Anthropic: ${err.message}`,
+            },
+          }),
+        );
+      } else {
+        res.end();
+      }
+      finalize();
+      resolve();
+    });
+
+    upstreamReq.write(outboundBytes);
+    upstreamReq.end();
+  });
+}

package/src/proxy/sse.ts

@@ -0,0 +1,58 @@
+import type { SSEEvent } from "../types.js";
+
+/**
+ * Streaming SSE parser. Accepts raw bytes (UTF-8) progressively and yields
+ * complete events. Events are buffered until a blank-line terminator is seen.
+ *
+ * Anthropic's SSE format emits lines of the form:
+ *   event: message_start
+ *   data: {"type":"message_start", ...}
+ *
+ * separated by blank lines. We preserve unrecognized fields and pass raw
+ * bytes through unchanged so the downstream client sees a byte-faithful
+ * stream — we only parse a copy for accounting.
+ */
+export class SSEParser {
+  private buffer = "";
+  private eventName = "";
+  private dataLines: string[] = [];
+
+  push(chunk: string): SSEEvent[] {
+    this.buffer += chunk;
+    const events: SSEEvent[] = [];
+    let idx: number;
+    while ((idx = this.buffer.indexOf("\n")) !== -1) {
+      const line = this.buffer.slice(0, idx).replace(/\r$/, "");
+      this.buffer = this.buffer.slice(idx + 1);
+      if (line === "") {
+        if (this.dataLines.length > 0 || this.eventName !== "") {
+          events.push({
+            event: this.eventName || "message",
+            data: this.dataLines.join("\n"),
+          });
+        }
+        this.eventName = "";
+        this.dataLines = [];
+      } else if (line.startsWith(":")) {
+        // comment / keep-alive
+      } else if (line.startsWith("event:")) {
+        this.eventName = line.slice(6).trimStart();
+      } else if (line.startsWith("data:")) {
+        this.dataLines.push(line.slice(5).trimStart());
+      }
+      // ignore other field names (id:, retry:) — not used by Anthropic
+    }
+    return events;
+  }
+
+  flush(): SSEEvent[] {
+    if (this.dataLines.length === 0 && this.eventName === "") return [];
+    const event = {
+      event: this.eventName || "message",
+      data: this.dataLines.join("\n"),
+    };
+    this.eventName = "";
+    this.dataLines = [];
+    return [event];
+  }
+}

package/src/proxy/usage.ts

@@ -0,0 +1,98 @@
+import type { SSEEvent, UsageCounts } from "../types.js";
+import { emptyUsage, addUsage } from "../pricing.js";
+
+interface AnthropicUsage {
+  input_tokens?: number;
+  output_tokens?: number;
+  cache_creation_input_tokens?: number;
+  cache_read_input_tokens?: number;
+}
+
+function fromAnthropic(u: AnthropicUsage | undefined): UsageCounts {
+  if (!u) return emptyUsage();
+  return {
+    inputTokens: u.input_tokens ?? 0,
+    outputTokens: u.output_tokens ?? 0,
+    cacheCreationInputTokens: u.cache_creation_input_tokens ?? 0,
+    cacheReadInputTokens: u.cache_read_input_tokens ?? 0,
+  };
+}
+
+/**
+ * Accumulates Anthropic usage from a stream of SSE events.
+ *
+ * Streaming responses emit:
+ *   message_start  — has usage with input_tokens + initial output_tokens=1
+ *   message_delta  — has usage with cumulative output_tokens
+ *   message_stop   — terminal
+ *
+ * Non-streaming responses come as a single JSON body with `.usage` at the top
+ * level — handled by usageFromJson.
+ */
+export class StreamUsageAccumulator {
+  private current: UsageCounts = emptyUsage();
+  private modelFromEvent: string | null = null;
+
+  observe(event: SSEEvent): void {
+    if (event.event !== "message_start" && event.event !== "message_delta") {
+      return;
+    }
+    let parsed: unknown;
+    try {
+      parsed = JSON.parse(event.data);
+    } catch {
+      return;
+    }
+    if (!parsed || typeof parsed !== "object") return;
+    const obj = parsed as Record<string, unknown>;
+
+    if (event.event === "message_start") {
+      const message = obj["message"] as Record<string, unknown> | undefined;
+      if (message) {
+        if (typeof message["model"] === "string") {
+          this.modelFromEvent = message["model"] as string;
+        }
+        const u = fromAnthropic(message["usage"] as AnthropicUsage | undefined);
+        this.current = u;
+      }
+    } else if (event.event === "message_delta") {
+      const usage = obj["usage"] as AnthropicUsage | undefined;
+      if (usage) {
+        // message_delta usage is cumulative for output_tokens; input is fixed
+        const u = fromAnthropic(usage);
+        this.current = {
+          inputTokens: this.current.inputTokens || u.inputTokens,
+          outputTokens: u.outputTokens,
+          cacheCreationInputTokens:
+            this.current.cacheCreationInputTokens ||
+            u.cacheCreationInputTokens,
+          cacheReadInputTokens:
+            this.current.cacheReadInputTokens || u.cacheReadInputTokens,
+        };
+      }
+    }
+  }
+
+  total(): UsageCounts {
+    return { ...this.current };
+  }
+
+  model(): string | null {
+    return this.modelFromEvent;
+  }
+}
+
+export function usageFromJson(body: unknown): {
+  usage: UsageCounts;
+  model: string | null;
+} {
+  if (!body || typeof body !== "object") {
+    return { usage: emptyUsage(), model: null };
+  }
+  const obj = body as Record<string, unknown>;
+  const usage = fromAnthropic(obj["usage"] as AnthropicUsage | undefined);
+  const model = typeof obj["model"] === "string" ? (obj["model"] as string) : null;
+  return { usage, model };
+}
+
+export { addUsage };

package/src/server.ts

@@ -0,0 +1,154 @@
+import { createServer, IncomingMessage, ServerResponse, Server } from "node:http";
+import type { ProxyConfig, RequestRecord } from "./types.js";
+import { handleAnthropicRequest } from "./proxy/anthropic-passthrough.js";
+import { Ledger } from "./ledger.js";
+
+export interface ProxyServerHandle {
+  proxy: Server;
+  dashboard: Server;
+  ledger: Ledger;
+  close: () => Promise<void>;
+}
+
+type DashboardRenderer = (ledger: Ledger) => string;
+
+export interface StartOptions {
+  config: ProxyConfig;
+  onRecord?: (r: RequestRecord) => void;
+  renderDashboard?: DashboardRenderer;
+}
+
+export function defaultConfig(overrides: Partial<ProxyConfig> = {}): ProxyConfig {
+  const home = process.env["HOME"] ?? process.env["USERPROFILE"] ?? ".";
+  return {
+    upstreamBaseUrl: overrides.upstreamBaseUrl ?? "https://api.anthropic.com",
+    port: overrides.port ?? 7777,
+    bind: overrides.bind ?? "127.0.0.1",
+    dashboardPort: overrides.dashboardPort ?? 7778,
+    ledgerPath: overrides.ledgerPath ?? `${home}/.tokenshield/ledger.db`,
+    enabledProcessors: overrides.enabledProcessors ?? ["token-accounting"],
+    retentionDays: overrides.retentionDays ?? 7,
+  };
+}
+
+async function listenOn(server: Server, port: number, bind: string): Promise<void> {
+  return new Promise((resolve, reject) => {
+    server.once("error", reject);
+    server.listen(port, bind, () => {
+      server.removeListener("error", reject);
+      resolve();
+    });
+  });
+}
+
+async function closeServer(server: Server): Promise<void> {
+  return new Promise((resolve) => {
+    server.close(() => resolve());
+  });
+}
+
+export async function start(opts: StartOptions): Promise<ProxyServerHandle> {
+  const ledger = new Ledger(opts.config.ledgerPath);
+
+  const sink = (r: RequestRecord): void => {
+    try {
+      ledger.record(r);
+    } catch {
+      // ledger errors must never break the request path
+    }
+    opts.onRecord?.(r);
+  };
+
+  const proxy = createServer((req: IncomingMessage, res: ServerResponse) => {
+    if (req.url === "/__tokenshield/health") {
+      res.statusCode = 200;
+      res.setHeader("content-type", "application/json");
+      res.end(JSON.stringify({ ok: true, version: "0.1.0" }));
+      return;
+    }
+    handleAnthropicRequest(req, res, opts.config, sink).catch((err: unknown) => {
+      if (!res.headersSent) {
+        res.statusCode = 500;
+        res.setHeader("content-type", "application/json");
+        res.end(
+          JSON.stringify({
+            type: "error",
+            error: {
+              type: "tokenshield_internal_error",
+              message: (err as Error)?.message ?? "unknown",
+            },
+          }),
+        );
+      } else {
+        try {
+          res.end();
+        } catch {
+          // ignore
+        }
+      }
+    });
+  });
+  proxy.keepAliveTimeout = 65_000;
+  proxy.headersTimeout = 70_000;
+  proxy.requestTimeout = 0; // streaming responses can be long
+
+  const dashboard = createServer((req: IncomingMessage, res: ServerResponse) => {
+    const url = req.url ?? "/";
+    if (url === "/api/summary") {
+      const since = Date.now() - 24 * 60 * 60 * 1000;
+      const summary = ledger.summary(since);
+      res.statusCode = 200;
+      res.setHeader("content-type", "application/json");
+      res.setHeader("cache-control", "no-store");
+      res.end(JSON.stringify(summary));
+      return;
+    }
+    if (url === "/api/recent") {
+      res.statusCode = 200;
+      res.setHeader("content-type", "application/json");
+      res.setHeader("cache-control", "no-store");
+      res.end(JSON.stringify(ledger.recent(50)));
+      return;
+    }
+    if (url === "/health") {
+      res.statusCode = 200;
+      res.setHeader("content-type", "application/json");
+      res.end(JSON.stringify({ ok: true }));
+      return;
+    }
+    const html = opts.renderDashboard?.(ledger) ?? defaultDashboardHtml();
+    res.statusCode = 200;
+    res.setHeader("content-type", "text/html; charset=utf-8");
+    res.setHeader("cache-control", "no-store");
+    res.end(html);
+  });
+
+  await listenOn(proxy, opts.config.port, opts.config.bind);
+  await listenOn(dashboard, opts.config.dashboardPort, opts.config.bind);
+
+  const retentionInterval = setInterval(() => {
+    const cutoff = Date.now() - opts.config.retentionDays * 24 * 60 * 60 * 1000;
+    try {
+      ledger.prune(cutoff);
+    } catch {
+      // ignore
+    }
+  }, 60 * 60 * 1000);
+  retentionInterval.unref();
+
+  return {
+    proxy,
+    dashboard,
+    ledger,
+    close: async () => {
+      clearInterval(retentionInterval);
+      await Promise.all([closeServer(proxy), closeServer(dashboard)]);
+      ledger.close();
+    },
+  };
+}
+
+function defaultDashboardHtml(): string {
+  return `<!doctype html><meta charset="utf-8"><title>TokenShield</title>
+<body><h1>TokenShield</h1><p>Dashboard renderer not provided.</p></body>`;
+}