@mcp-audit-gateway/core 0.1.0

package/src/proxy/upstream-manager.ts

@@ -0,0 +1,326 @@
+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
+import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
+import type { UpstreamConfig, UpstreamStatus } from "../types.js";
+
+export interface UpstreamConnection {
+  config: UpstreamConfig;
+  client: Client;
+  status: UpstreamStatus;
+  tools: Array<{
+    name: string;
+    description?: string;
+    inputSchema?: unknown;
+    annotations?: Record<string, unknown>;
+  }>;
+  consecutiveFailures: number;
+  reconnectAttempts: number;
+  reconnectTimer?: ReturnType<typeof setTimeout>;
+}
+
+export class UpstreamManager {
+  private connections: Map<string, UpstreamConnection> = new Map();
+  private healthCheckInterval?: ReturnType<typeof setInterval>;
+  private static readonly MAX_CONSECUTIVE_FAILURES = 3;
+  private static readonly DEFAULT_TIMEOUT_MS = 10_000;
+  private static readonly MAX_BACKOFF_MS = 30_000;
+  private static readonly BASE_BACKOFF_MS = 1_000;
+
+  async connect(
+    upstream: UpstreamConfig,
+    timeoutMs: number = UpstreamManager.DEFAULT_TIMEOUT_MS,
+  ): Promise<UpstreamConnection> {
+    const client = new Client(
+      { name: "mcp-audit-gateway", version: "0.1.0" },
+      { capabilities: {} },
+    );
+
+    const transport = this.createTransport(upstream);
+
+    await this.withTimeout(client.connect(transport), timeoutMs, upstream.name);
+
+    let tools: UpstreamConnection["tools"];
+    try {
+      const toolsResult = await this.withTimeout(
+        client.listTools(),
+        timeoutMs,
+        upstream.name,
+      );
+      tools = toolsResult.tools.map((t) => ({
+        name: t.name,
+        description: t.description,
+        inputSchema: t.inputSchema,
+        annotations: t.annotations as Record<string, unknown> | undefined,
+      }));
+    } catch (err) {
+      try {
+        await client.close();
+      } catch {
+        // best effort
+      }
+      throw err;
+    }
+
+    const connection: UpstreamConnection = {
+      config: upstream,
+      client,
+      status: {
+        name: upstream.name,
+        namespace: upstream.namespace,
+        status: "healthy",
+        toolCount: tools.length,
+        lastSuccessfulContact: new Date().toISOString(),
+      },
+      tools,
+      consecutiveFailures: 0,
+      reconnectAttempts: 0,
+    };
+
+    this.connections.set(upstream.name, connection);
+    return connection;
+  }
+
+  async callTool(
+    upstreamName: string,
+    toolName: string,
+    args: Record<string, unknown>,
+  ): Promise<unknown> {
+    const conn = this.connections.get(upstreamName);
+    if (!conn) {
+      throw new Error(`Upstream not connected: ${upstreamName}`);
+    }
+
+    if (conn.status.status === "unavailable") {
+      throw new UpstreamUnavailableError(
+        upstreamName,
+        conn.status.unavailableReason ?? "Upstream is unavailable",
+      );
+    }
+
+    try {
+      const result = await conn.client.callTool({ name: toolName, arguments: args });
+      conn.status.lastSuccessfulContact = new Date().toISOString();
+      conn.status.status = "healthy";
+      conn.consecutiveFailures = 0;
+      return result;
+    } catch (err) {
+      conn.consecutiveFailures++;
+      if (conn.consecutiveFailures >= UpstreamManager.MAX_CONSECUTIVE_FAILURES) {
+        this.markUnavailable(conn, String(err));
+      } else {
+        conn.status.status = "degraded";
+        conn.status.degradedReason = String(err);
+      }
+      throw err;
+    }
+  }
+
+  startHealthChecks(intervalMs: number): void {
+    this.stopHealthChecks();
+    this.healthCheckInterval = setInterval(() => {
+      void this.runHealthChecks();
+    }, intervalMs);
+  }
+
+  stopHealthChecks(): void {
+    if (this.healthCheckInterval) {
+      clearInterval(this.healthCheckInterval);
+      this.healthCheckInterval = undefined;
+    }
+  }
+
+  async reconnect(upstreamName: string): Promise<void> {
+    const conn = this.connections.get(upstreamName);
+    if (!conn) {
+      throw new Error(`Upstream not found: ${upstreamName}`);
+    }
+
+    if (conn.reconnectTimer) {
+      clearTimeout(conn.reconnectTimer);
+      conn.reconnectTimer = undefined;
+    }
+
+    try {
+      try {
+        await conn.client.close();
+      } catch {
+        // best effort
+      }
+
+      const client = new Client(
+        { name: "mcp-audit-gateway", version: "0.1.0" },
+        { capabilities: {} },
+      );
+      const transport = this.createTransport(conn.config);
+      await this.withTimeout(
+        client.connect(transport),
+        UpstreamManager.DEFAULT_TIMEOUT_MS,
+        upstreamName,
+      );
+
+      const toolsResult = await this.withTimeout(
+        client.listTools(),
+        UpstreamManager.DEFAULT_TIMEOUT_MS,
+        upstreamName,
+      );
+      const tools = toolsResult.tools.map((t) => ({
+        name: t.name,
+        description: t.description,
+        inputSchema: t.inputSchema,
+        annotations: t.annotations as Record<string, unknown> | undefined,
+      }));
+
+      conn.client = client;
+      conn.tools = tools;
+      conn.status.status = "healthy";
+      conn.status.toolCount = tools.length;
+      conn.status.lastSuccessfulContact = new Date().toISOString();
+      conn.status.unavailableReason = undefined;
+      conn.status.degradedReason = undefined;
+      conn.consecutiveFailures = 0;
+      conn.reconnectAttempts = 0;
+    } catch (err) {
+      conn.reconnectAttempts++;
+      this.scheduleReconnect(conn);
+      throw err;
+    }
+  }
+
+  getConnection(name: string): UpstreamConnection | undefined {
+    return this.connections.get(name);
+  }
+
+  getAllStatuses(): UpstreamStatus[] {
+    return Array.from(this.connections.values()).map((c) => c.status);
+  }
+
+  async disconnectAll(): Promise<void> {
+    this.stopHealthChecks();
+    for (const conn of this.connections.values()) {
+      if (conn.reconnectTimer) {
+        clearTimeout(conn.reconnectTimer);
+      }
+      try {
+        await conn.client.close();
+      } catch {
+        // best effort
+      }
+    }
+    this.connections.clear();
+  }
+
+  private async runHealthChecks(): Promise<void> {
+    for (const conn of this.connections.values()) {
+      if (conn.status.status === "unavailable") {
+        continue;
+      }
+
+      try {
+        await conn.client.ping();
+        conn.status.lastSuccessfulContact = new Date().toISOString();
+        conn.consecutiveFailures = 0;
+        if (conn.status.status === "degraded") {
+          conn.status.status = "healthy";
+          conn.status.degradedReason = undefined;
+        }
+      } catch (err) {
+        conn.consecutiveFailures++;
+        if (conn.consecutiveFailures >= UpstreamManager.MAX_CONSECUTIVE_FAILURES) {
+          this.markUnavailable(conn, `Health check failed: ${err}`);
+        } else {
+          conn.status.status = "degraded";
+          conn.status.degradedReason = `Health check failed: ${err}`;
+        }
+      }
+    }
+  }
+
+  private markUnavailable(conn: UpstreamConnection, reason: string): void {
+    conn.status.status = "unavailable";
+    conn.status.unavailableReason = reason;
+    conn.status.degradedReason = undefined;
+    this.scheduleReconnect(conn);
+  }
+
+  private scheduleReconnect(conn: UpstreamConnection): void {
+    if (conn.reconnectTimer) {
+      clearTimeout(conn.reconnectTimer);
+    }
+
+    const backoffMs = this.calculateBackoff(conn.reconnectAttempts);
+    conn.reconnectTimer = setTimeout(() => {
+      void this.reconnect(conn.config.name).catch(() => {});
+    }, backoffMs);
+  }
+
+  private calculateBackoff(attempt: number): number {
+    const exponential = UpstreamManager.BASE_BACKOFF_MS * Math.pow(2, attempt);
+    return Math.min(exponential, UpstreamManager.MAX_BACKOFF_MS);
+  }
+
+  private async withTimeout<T>(
+    promise: Promise<T>,
+    timeoutMs: number,
+    upstreamName: string,
+  ): Promise<T> {
+    let timer: ReturnType<typeof setTimeout>;
+    const timeoutPromise = new Promise<never>((_, reject) => {
+      timer = setTimeout(() => {
+        reject(
+          new ConnectionTimeoutError(
+            upstreamName,
+            timeoutMs,
+          ),
+        );
+      }, timeoutMs);
+    });
+
+    try {
+      const result = await Promise.race([promise, timeoutPromise]);
+      clearTimeout(timer!);
+      return result;
+    } catch (err) {
+      clearTimeout(timer!);
+      throw err;
+    }
+  }
+
+  private createTransport(upstream: UpstreamConfig) {
+    if (upstream.transport.type === "stdio") {
+      return new StdioClientTransport({
+        command: upstream.transport.command,
+        args: upstream.transport.args,
+        env: upstream.transport.env,
+      });
+    }
+
+    return new StreamableHTTPClientTransport(
+      new URL(upstream.transport.url),
+      upstream.transport.headers ? { requestInit: { headers: upstream.transport.headers } } : undefined,
+    );
+  }
+}
+
+export class UpstreamUnavailableError extends Error {
+  constructor(
+    public readonly upstreamName: string,
+    public readonly reason: string,
+  ) {
+    super(
+      `Upstream "${upstreamName}" is unavailable: ${reason}`,
+    );
+    this.name = "UpstreamUnavailableError";
+  }
+}
+
+export class ConnectionTimeoutError extends Error {
+  constructor(
+    public readonly upstreamName: string,
+    public readonly timeoutMs: number,
+  ) {
+    super(
+      `Connection to upstream "${upstreamName}" timed out after ${timeoutMs}ms`,
+    );
+    this.name = "ConnectionTimeoutError";
+  }
+}

package/src/telemetry/logger.test.ts

@@ -0,0 +1,139 @@
+import { describe, it, expect, beforeEach, vi } from "vitest";
+import { Logger } from "./logger.js";
+import type { LogLevel } from "./logger.js";
+
+describe("Logger", () => {
+  let lines: string[];
+  let logger: Logger;
+
+  beforeEach(() => {
+    lines = [];
+    logger = new Logger({
+      component: "test",
+      level: "debug",
+      output: (line) => lines.push(line),
+    });
+  });
+
+  describe("structured output", () => {
+    it("outputs valid JSON lines", () => {
+      logger.info("hello world");
+      expect(lines).toHaveLength(1);
+      const parsed = JSON.parse(lines[0]);
+      expect(parsed).toHaveProperty("timestamp");
+      expect(parsed).toHaveProperty("level", "info");
+      expect(parsed).toHaveProperty("component", "test");
+      expect(parsed).toHaveProperty("message", "hello world");
+    });
+
+    it("includes ISO 8601 timestamp", () => {
+      logger.info("msg");
+      const parsed = JSON.parse(lines[0]);
+      expect(() => new Date(parsed.timestamp)).not.toThrow();
+      expect(parsed.timestamp).toMatch(/^\d{4}-\d{2}-\d{2}T/);
+    });
+
+    it("includes arbitrary context fields", () => {
+      logger.info("tool call", { tool: "my_tool", durationMs: 42, upstream: "server-a" });
+      const parsed = JSON.parse(lines[0]);
+      expect(parsed.tool).toBe("my_tool");
+      expect(parsed.durationMs).toBe(42);
+      expect(parsed.upstream).toBe("server-a");
+    });
+  });
+
+  describe("log levels", () => {
+    it("outputs debug messages when level is debug", () => {
+      logger.debug("trace msg");
+      expect(lines).toHaveLength(1);
+      expect(JSON.parse(lines[0]).level).toBe("debug");
+    });
+
+    it("outputs info messages", () => {
+      logger.info("info msg");
+      expect(lines).toHaveLength(1);
+      expect(JSON.parse(lines[0]).level).toBe("info");
+    });
+
+    it("outputs warn messages", () => {
+      logger.warn("warning");
+      expect(lines).toHaveLength(1);
+      expect(JSON.parse(lines[0]).level).toBe("warn");
+    });
+
+    it("outputs error messages", () => {
+      logger.error("failure");
+      expect(lines).toHaveLength(1);
+      expect(JSON.parse(lines[0]).level).toBe("error");
+    });
+  });
+
+  describe("level filtering", () => {
+    it("suppresses debug when level is info", () => {
+      const infoLogger = new Logger({
+        component: "test",
+        level: "info",
+        output: (line) => lines.push(line),
+      });
+      infoLogger.debug("should not appear");
+      infoLogger.info("should appear");
+      expect(lines).toHaveLength(1);
+      expect(JSON.parse(lines[0]).message).toBe("should appear");
+    });
+
+    it("suppresses info and debug when level is warn", () => {
+      const warnLogger = new Logger({
+        component: "test",
+        level: "warn",
+        output: (line) => lines.push(line),
+      });
+      warnLogger.debug("no");
+      warnLogger.info("no");
+      warnLogger.warn("yes");
+      warnLogger.error("yes");
+      expect(lines).toHaveLength(2);
+    });
+
+    it("only outputs error when level is error", () => {
+      const errorLogger = new Logger({
+        component: "test",
+        level: "error",
+        output: (line) => lines.push(line),
+      });
+      errorLogger.debug("no");
+      errorLogger.info("no");
+      errorLogger.warn("no");
+      errorLogger.error("yes");
+      expect(lines).toHaveLength(1);
+    });
+  });
+
+  describe("child logger", () => {
+    it("creates child with different component", () => {
+      const child = logger.child({ component: "sub-system" });
+      child.info("child message");
+      const parsed = JSON.parse(lines[0]);
+      expect(parsed.component).toBe("sub-system");
+    });
+
+    it("child inherits parent output", () => {
+      const child = logger.child({ component: "child" });
+      child.warn("from child");
+      expect(lines).toHaveLength(1);
+      expect(JSON.parse(lines[0]).component).toBe("child");
+    });
+  });
+
+  describe("default output", () => {
+    it("writes to stderr by default", () => {
+      const stderrWrite = vi.spyOn(process.stderr, "write").mockImplementation(() => true);
+      const defaultLogger = new Logger({ component: "default-test", level: "info" });
+      defaultLogger.info("stderr test");
+      expect(stderrWrite).toHaveBeenCalled();
+      const written = stderrWrite.mock.calls[0][0] as string;
+      expect(written).toContain("stderr test");
+      expect(written.endsWith("\n")).toBe(true);
+      stderrWrite.mockRestore();
+    });
+  });
+});

package/src/telemetry/logger.ts

@@ -0,0 +1,72 @@
+export type LogLevel = "debug" | "info" | "warn" | "error";
+
+const LEVEL_PRIORITY: Record<LogLevel, number> = {
+  debug: 0,
+  info: 1,
+  warn: 2,
+  error: 3,
+};
+
+export interface LogEntry {
+  timestamp: string;
+  level: LogLevel;
+  component: string;
+  message: string;
+  [key: string]: unknown;
+}
+
+export interface LoggerOptions {
+  component: string;
+  level?: LogLevel;
+  output?: (line: string) => void;
+}
+
+export class Logger {
+  private component: string;
+  private minLevel: number;
+  private output: (line: string) => void;
+
+  constructor(options: LoggerOptions) {
+    this.component = options.component;
+    this.minLevel = LEVEL_PRIORITY[options.level ?? "info"];
+    this.output = options.output ?? ((line: string) => process.stderr.write(line + "\n"));
+  }
+
+  debug(message: string, context?: Record<string, unknown>): void {
+    this.log("debug", message, context);
+  }
+
+  info(message: string, context?: Record<string, unknown>): void {
+    this.log("info", message, context);
+  }
+
+  warn(message: string, context?: Record<string, unknown>): void {
+    this.log("warn", message, context);
+  }
+
+  error(message: string, context?: Record<string, unknown>): void {
+    this.log("error", message, context);
+  }
+
+  private log(level: LogLevel, message: string, context?: Record<string, unknown>): void {
+    if (LEVEL_PRIORITY[level] < this.minLevel) return;
+
+    const entry: LogEntry = {
+      timestamp: new Date().toISOString(),
+      level,
+      component: this.component,
+      message,
+      ...context,
+    };
+
+    this.output(JSON.stringify(entry));
+  }
+
+  child(overrides: Partial<LoggerOptions>): Logger {
+    return new Logger({
+      component: overrides.component ?? this.component,
+      level: overrides.level ?? (Object.keys(LEVEL_PRIORITY) as LogLevel[])[this.minLevel],
+      output: overrides.output ?? this.output,
+    });
+  }
+}

package/src/telemetry/metrics.test.ts

@@ -0,0 +1,125 @@
+import { describe, it, expect, beforeEach, vi } from "vitest";
+import { metrics } from "@opentelemetry/api";
+import { GatewayMetrics } from "./metrics.js";
+import type { TelemetryConfig } from "../types.js";
+
+const testConfig: TelemetryConfig = {
+  enabled: true,
+  serviceName: "test-gateway",
+  sampleRate: 1.0,
+};
+
+describe("GatewayMetrics", () => {
+  let gatewayMetrics: GatewayMetrics;
+
+  beforeEach(() => {
+    gatewayMetrics = new GatewayMetrics(testConfig);
+  });
+
+  describe("construction", () => {
+    it("creates all metric instruments", () => {
+      expect(gatewayMetrics.toolCallsTotal).toBeDefined();
+      expect(gatewayMetrics.toolCallDurationMs).toBeDefined();
+      expect(gatewayMetrics.upstreamStatus).toBeDefined();
+      expect(gatewayMetrics.policyDenialsTotal).toBeDefined();
+    });
+  });
+
+  describe("recordToolCall", () => {
+    it("increments tool_calls_total counter without throwing", () => {
+      expect(() =>
+        gatewayMetrics.recordToolCall({
+          namespace: "test",
+          tool: "test/my_tool",
+          principal: "agent:user1",
+          success: "true",
+        }),
+      ).not.toThrow();
+    });
+
+    it("accepts different label combinations", () => {
+      expect(() =>
+        gatewayMetrics.recordToolCall({
+          namespace: "prod",
+          tool: "prod/fetch_data",
+          principal: "anonymous",
+          success: "false",
+        }),
+      ).not.toThrow();
+    });
+  });
+
+  describe("recordDuration", () => {
+    it("records duration histogram without throwing", () => {
+      expect(() =>
+        gatewayMetrics.recordDuration(42.5, {
+          namespace: "test",
+          upstream: "my-server",
+        }),
+      ).not.toThrow();
+    });
+
+    it("handles zero duration", () => {
+      expect(() =>
+        gatewayMetrics.recordDuration(0, {
+          namespace: "test",
+          upstream: "fast-server",
+        }),
+      ).not.toThrow();
+    });
+
+    it("handles large durations", () => {
+      expect(() =>
+        gatewayMetrics.recordDuration(30000, {
+          namespace: "test",
+          upstream: "slow-server",
+        }),
+      ).not.toThrow();
+    });
+  });
+
+  describe("setUpstreamStatus", () => {
+    it("records gauge value for healthy status", () => {
+      expect(() =>
+        gatewayMetrics.setUpstreamStatus("my-upstream", "healthy", 1),
+      ).not.toThrow();
+    });
+
+    it("records gauge value for unhealthy status", () => {
+      expect(() =>
+        gatewayMetrics.setUpstreamStatus("my-upstream", "unavailable", 0),
+      ).not.toThrow();
+    });
+  });
+
+  describe("recordPolicyDenial", () => {
+    it("increments policy_denials_total counter without throwing", () => {
+      expect(() =>
+        gatewayMetrics.recordPolicyDenial({
+          principal: "agent:blocked",
+          tool: "test/dangerous_tool",
+          reason: "unauthorized",
+        }),
+      ).not.toThrow();
+    });
+
+    it("records denial with rate limit reason", () => {
+      expect(() =>
+        gatewayMetrics.recordPolicyDenial({
+          principal: "agent:heavy",
+          tool: "test/expensive_tool",
+          reason: "rate_limit_exceeded",
+        }),
+      ).not.toThrow();
+    });
+  });
+
+  describe("with custom MeterProvider", () => {
+    it("uses the global meter provider", () => {
+      const getMeterSpy = vi.spyOn(metrics, "getMeter");
+      new GatewayMetrics(testConfig);
+      expect(getMeterSpy).toHaveBeenCalledWith("mcp-audit-gateway", "0.1.0");
+      getMeterSpy.mockRestore();
+    });
+  });
+});