@langwatch/mcp-server 1.0.0 → 2.0.0

package/src/__tests__/http-security-sessions.unit.test.ts

@@ -0,0 +1,209 @@
+import { describe, expect, it, vi } from "vitest";
+
+import { admitOAuthToken, createSessionStore } from "../http-security.js";
+
+describe("createSessionStore", () => {
+  function makeStore(maxAgeMs = 60_000) {
+    const closed: string[] = [];
+    const store = createSessionStore<{ id: string }>({
+      maxAgeMs,
+      closeTransport: (transport) => closed.push(transport.id),
+    });
+    return { store, closed };
+  }
+
+  it("stores and returns sessions bound to their key", () => {
+    const { store } = makeStore();
+    store.add({ sessionId: "s1", transport: { id: "s1" }, apiKey: "sk-a" });
+
+    expect(store.get("s1")?.apiKey).toBe("sk-a");
+    expect(store.size).toBe(1);
+  });
+
+  it("counts live sessions per key without holding the raw key", () => {
+    const { store } = makeStore();
+    store.add({ sessionId: "s1", transport: { id: "s1" }, apiKey: "sk-a" });
+    store.add({ sessionId: "s2", transport: { id: "s2" }, apiKey: "sk-a" });
+    store.add({ sessionId: "s3", transport: { id: "s3" }, apiKey: "sk-b" });
+
+    expect(store.countForKey("sk-a")).toBe(2);
+    expect(store.countForKey("sk-b")).toBe(1);
+    expect(store.countForKey("sk-unknown")).toBe(0);
+  });
+
+  it("decrements the per-key count when a session is removed", () => {
+    const { store } = makeStore();
+    store.add({ sessionId: "s1", transport: { id: "s1" }, apiKey: "sk-a" });
+    store.add({ sessionId: "s2", transport: { id: "s2" }, apiKey: "sk-a" });
+
+    store.remove("s1");
+    expect(store.countForKey("sk-a")).toBe(1);
+
+    store.remove("s2");
+    expect(store.countForKey("sk-a")).toBe(0);
+  });
+
+  it("ignores removal of a session it does not hold", () => {
+    const { store } = makeStore();
+    store.add({ sessionId: "s1", transport: { id: "s1" }, apiKey: "sk-a" });
+
+    store.remove("never-existed");
+    store.remove("never-existed");
+
+    expect(store.countForKey("sk-a")).toBe(1);
+  });
+
+  it("closes and forgets sessions that have gone idle", () => {
+    vi.useFakeTimers();
+    try {
+      const { store, closed } = makeStore(1_000);
+      store.add({ sessionId: "idle", transport: { id: "idle" }, apiKey: "sk-a" });
+      store.add({ sessionId: "busy", transport: { id: "busy" }, apiKey: "sk-a" });
+
+      vi.advanceTimersByTime(900);
+      store.touch("busy");
+      vi.advanceTimersByTime(900);
+      store.sweep();
+
+      expect(closed).toEqual(["idle"]);
+      expect(store.get("idle")).toBeUndefined();
+      expect(store.get("busy")).toBeDefined();
+      expect(store.countForKey("sk-a")).toBe(1);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("closes everything on shutdown", () => {
+    const { store, closed } = makeStore();
+    store.add({ sessionId: "s1", transport: { id: "s1" }, apiKey: "sk-a" });
+    store.add({ sessionId: "s2", transport: { id: "s2" }, apiKey: "sk-b" });
+
+    store.closeAll();
+
+    expect(closed.sort()).toEqual(["s1", "s2"]);
+    expect(store.size).toBe(0);
+    expect(store.countForKey("sk-a")).toBe(0);
+  });
+
+  it("does not let a prototype-shaped session id reach the prototype chain", () => {
+    const { store } = makeStore();
+
+    expect(store.get("__proto__")).toBeUndefined();
+    expect(store.get("constructor")).toBeUndefined();
+
+    store.add({ sessionId: "__proto__", transport: { id: "weird" }, apiKey: "sk-a" });
+    expect(store.get("__proto__")?.transport.id).toBe("weird");
+    expect(store.size).toBe(1);
+  });
+});
+
+describe("session reservations", () => {
+  function makeStore() {
+    return createSessionStore<{ id: string }>({
+      maxAgeMs: 60_000,
+      closeTransport: () => undefined,
+    });
+  }
+
+  it("counts a reserved slot before any session exists", () => {
+    const store = makeStore();
+    store.reserve("sk-a");
+
+    // The point of the reservation: the cap sees the slot while the session is
+    // still being set up.
+    expect(store.countForKey("sk-a")).toBe(1);
+    expect(store.size).toBe(0);
+  });
+
+  it("hands the reserved slot to the session rather than counting twice", () => {
+    const store = makeStore();
+    store.reserve("sk-a").commit({ sessionId: "s1", transport: { id: "s1" } });
+
+    expect(store.countForKey("sk-a")).toBe(1);
+    expect(store.get("s1")?.apiKey).toBe("sk-a");
+  });
+
+  it("returns the slot when the session never came up", () => {
+    const store = makeStore();
+    const reservation = store.reserve("sk-a");
+    reservation.release();
+
+    expect(store.countForKey("sk-a")).toBe(0);
+  });
+
+  it("ignores a release once the slot was committed", () => {
+    const store = makeStore();
+    const reservation = store.reserve("sk-a");
+    reservation.commit({ sessionId: "s1", transport: { id: "s1" } });
+    reservation.release();
+
+    expect(store.countForKey("sk-a")).toBe(1);
+    expect(store.get("s1")).toBeDefined();
+  });
+
+  it("frees the slot when the committed session is removed", () => {
+    const store = makeStore();
+    store.reserve("sk-a").commit({ sessionId: "s1", transport: { id: "s1" } });
+    store.remove("s1");
+
+    expect(store.countForKey("sk-a")).toBe(0);
+  });
+});
+
+describe("admitOAuthToken", () => {
+  function entry(apiKey: string, expiresAt: number) {
+    return { apiKey, expiresAt };
+  }
+
+  it("evicts the oldest token for a key that is at the cap", () => {
+    const tokens = new Map([
+      ["t1", entry("sk-a", Date.now() + 1_000)],
+      ["t2", entry("sk-a", Date.now() + 2_000)],
+      ["t3", entry("sk-a", Date.now() + 3_000)],
+    ]);
+
+    admitOAuthToken({ apiKey: "sk-a", tokens, maxPerKey: 3 });
+
+    // Room for the caller's new token, and the longest-standing one went.
+    expect(tokens.has("t1")).toBe(false);
+    expect([...tokens.keys()]).toEqual(["t2", "t3"]);
+  });
+
+  it("leaves other keys alone", () => {
+    const tokens = new Map([
+      ["t1", entry("sk-a", Date.now() + 1_000)],
+      ["t2", entry("sk-b", Date.now() + 1_000)],
+      ["t3", entry("sk-b", Date.now() + 2_000)],
+    ]);
+
+    admitOAuthToken({ apiKey: "sk-a", tokens, maxPerKey: 1 });
+
+    expect(tokens.has("t1")).toBe(false);
+    expect(tokens.has("t2")).toBe(true);
+    expect(tokens.has("t3")).toBe(true);
+  });
+
+  it("keeps a key from growing the map without bound", () => {
+    const tokens = new Map<string, { apiKey: string; expiresAt: number }>();
+
+    for (let i = 0; i < 50; i += 1) {
+      admitOAuthToken({ apiKey: "sk-a", tokens, maxPerKey: 10 });
+      tokens.set(`t${i}`, entry("sk-a", Date.now() + 60_000 + i));
+    }
+
+    expect(tokens.size).toBe(10);
+  });
+
+  it("drops expired tokens for any key on the way past", () => {
+    const tokens = new Map([
+      ["stale", entry("sk-b", Date.now() - 1)],
+      ["live", entry("sk-b", Date.now() + 60_000)],
+    ]);
+
+    admitOAuthToken({ apiKey: "sk-a", tokens, maxPerKey: 10 });
+
+    expect(tokens.has("stale")).toBe(false);
+    expect(tokens.has("live")).toBe(true);
+  });
+});

package/src/__tests__/http-security-verifier.unit.test.ts

@@ -0,0 +1,248 @@
+import { describe, expect, it, vi } from "vitest";
+
+import {
+  apiKeysMatch,
+  createApiKeyVerifier,
+  createRateLimiter,
+  hashApiKey,
+} from "../http-security.js";
+
+describe("apiKeysMatch", () => {
+  it("matches identical keys", () => {
+    expect(apiKeysMatch({ presentedKey: "sk-abc", expectedKey: "sk-abc" })).toBe(true);
+  });
+
+  it("rejects different keys, including different lengths", () => {
+    expect(apiKeysMatch({ presentedKey: "sk-abc", expectedKey: "sk-abd" })).toBe(false);
+    expect(apiKeysMatch({ presentedKey: "sk-abc", expectedKey: "sk-abc-longer" })).toBe(false);
+    expect(apiKeysMatch({ presentedKey: "", expectedKey: "sk-abc" })).toBe(false);
+  });
+});
+
+describe("hashApiKey", () => {
+  it("never returns the raw key", () => {
+    expect(hashApiKey("sk-secret")).not.toContain("sk-secret");
+  });
+
+  it("is stable and distinct per key", () => {
+    expect(hashApiKey("a")).toBe(hashApiKey("a"));
+    expect(hashApiKey("a")).not.toBe(hashApiKey("b"));
+  });
+});
+
+describe("createRateLimiter", () => {
+  it("blocks only after the limit is exceeded", () => {
+    const limiter = createRateLimiter({ windowMs: 60_000, maxRequests: 3 });
+
+    expect(limiter.isBlocked("1.2.3.4")).toBe(false);
+    limiter.track("1.2.3.4");
+    limiter.track("1.2.3.4");
+    expect(limiter.isBlocked("1.2.3.4")).toBe(false);
+    limiter.track("1.2.3.4");
+    expect(limiter.isBlocked("1.2.3.4")).toBe(true);
+  });
+
+  it("tracks each IP separately", () => {
+    const limiter = createRateLimiter({ windowMs: 60_000, maxRequests: 1 });
+
+    limiter.track("1.1.1.1");
+    expect(limiter.isBlocked("1.1.1.1")).toBe(true);
+    expect(limiter.isBlocked("2.2.2.2")).toBe(false);
+  });
+
+  it("forgets an IP once its window has passed", () => {
+    vi.useFakeTimers();
+    try {
+      const limiter = createRateLimiter({ windowMs: 1_000, maxRequests: 1 });
+      limiter.track("1.1.1.1");
+      expect(limiter.isBlocked("1.1.1.1")).toBe(true);
+
+      vi.advanceTimersByTime(1_500);
+      expect(limiter.isBlocked("1.1.1.1")).toBe(false);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+});
+
+function jsonResponse(status: number): Response {
+  return new Response(status === 200 ? "{}" : "", { status });
+}
+
+/** A fetch double that records the URL and init it was called with. */
+function fetchStub(respond: () => Promise<Response>) {
+  return vi.fn(async (_url: string, _init?: RequestInit) => respond());
+}
+
+describe("createApiKeyVerifier", () => {
+  it("asks the LangWatch identity endpoint with the key in a header", async () => {
+    const fetchImpl = fetchStub(async () => jsonResponse(200));
+    const verifier = createApiKeyVerifier({
+      endpoint: "https://app.langwatch.ai",
+      fetchImpl: fetchImpl as unknown as typeof fetch,
+    });
+
+    await expect(verifier.verify("sk-real")).resolves.toBe(true);
+
+    expect(fetchImpl).toHaveBeenCalledWith(
+      "https://app.langwatch.ai/api/me/project",
+      expect.objectContaining({
+        method: "GET",
+        headers: { "X-Auth-Token": "sk-real" },
+      })
+    );
+  });
+
+  it("never puts the key in the URL", async () => {
+    const fetchImpl = fetchStub(async () => jsonResponse(200));
+    const verifier = createApiKeyVerifier({
+      endpoint: "https://app.langwatch.ai",
+      fetchImpl: fetchImpl as unknown as typeof fetch,
+    });
+
+    await verifier.verify("sk-real");
+
+    expect(fetchImpl.mock.calls[0]?.[0]).not.toContain("sk-real");
+  });
+
+  it.each([401, 403])("rejects a key the API answers %i for", async (status) => {
+    const verifier = createApiKeyVerifier({
+      endpoint: "https://app.langwatch.ai",
+      fetchImpl: (async () => jsonResponse(status)) as unknown as typeof fetch,
+    });
+
+    await expect(verifier.verify("sk-fake")).resolves.toBe(false);
+  });
+
+  it("caches a positive answer instead of asking again", async () => {
+    const fetchImpl = vi.fn(async () => jsonResponse(200));
+    const verifier = createApiKeyVerifier({
+      endpoint: "https://app.langwatch.ai",
+      fetchImpl: fetchImpl as unknown as typeof fetch,
+    });
+
+    await verifier.verify("sk-real");
+    await verifier.verify("sk-real");
+    await verifier.verify("sk-real");
+
+    expect(fetchImpl).toHaveBeenCalledTimes(1);
+  });
+
+  it("caches a negative answer so key spraying does not reach the API", async () => {
+    const fetchImpl = vi.fn(async () => jsonResponse(401));
+    const verifier = createApiKeyVerifier({
+      endpoint: "https://app.langwatch.ai",
+      fetchImpl: fetchImpl as unknown as typeof fetch,
+    });
+
+    await verifier.verify("sk-fake");
+    await verifier.verify("sk-fake");
+
+    expect(fetchImpl).toHaveBeenCalledTimes(1);
+  });
+
+  it("re-asks once the cache entry has expired", async () => {
+    vi.useFakeTimers();
+    try {
+      const fetchImpl = fetchStub(async () => jsonResponse(200));
+      const verifier = createApiKeyVerifier({
+        endpoint: "https://app.langwatch.ai",
+        positiveTtlMs: 1_000,
+        fetchImpl: fetchImpl as unknown as typeof fetch,
+      });
+
+      await verifier.verify("sk-real");
+      vi.advanceTimersByTime(1_500);
+      await verifier.verify("sk-real");
+
+      expect(fetchImpl).toHaveBeenCalledTimes(2);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("collapses concurrent checks of the same key into one request", async () => {
+    const fetchImpl = vi.fn(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 10));
+      return jsonResponse(200);
+    });
+    const verifier = createApiKeyVerifier({
+      endpoint: "https://app.langwatch.ai",
+      fetchImpl: fetchImpl as unknown as typeof fetch,
+    });
+
+    const results = await Promise.all(
+      Array.from({ length: 25 }, () => verifier.verify("sk-real"))
+    );
+
+    expect(results.every((result) => result === true)).toBe(true);
+    expect(fetchImpl).toHaveBeenCalledTimes(1);
+  });
+
+  it("fails closed when the API cannot answer, without caching the failure", async () => {
+    let status = 503;
+    const fetchImpl = vi.fn(async () => jsonResponse(status));
+    const verifier = createApiKeyVerifier({
+      endpoint: "https://app.langwatch.ai",
+      fetchImpl: fetchImpl as unknown as typeof fetch,
+    });
+
+    await expect(verifier.verify("sk-real")).resolves.toBe(false);
+
+    status = 200;
+    await expect(verifier.verify("sk-real")).resolves.toBe(true);
+    expect(fetchImpl).toHaveBeenCalledTimes(2);
+  });
+
+  it("gives up on an upstream that never answers", async () => {
+    // Without a timeout the in-flight promise never settles and every request
+    // for this key parks behind it until the socket dies.
+    const verifier = createApiKeyVerifier({
+      endpoint: "https://app.langwatch.ai",
+      requestTimeoutMs: 20,
+      fetchImpl: ((_url: string, init?: RequestInit) =>
+        new Promise((_resolve, reject) => {
+          init?.signal?.addEventListener("abort", () =>
+            reject(new Error("aborted"))
+          );
+        })) as unknown as typeof fetch,
+    });
+
+    await expect(verifier.verify("sk-real")).resolves.toBe(false);
+  });
+
+  it("fails closed when the request throws", async () => {
+    const verifier = createApiKeyVerifier({
+      endpoint: "https://app.langwatch.ai",
+      fetchImpl: (async () => {
+        throw new Error("connection refused");
+      }) as unknown as typeof fetch,
+    });
+
+    await expect(verifier.verify("sk-real")).resolves.toBe(false);
+  });
+
+  it("evicts old entries rather than growing once maxEntries is reached", async () => {
+    const fetchImpl = fetchStub(async () => jsonResponse(401));
+    const verifier = createApiKeyVerifier({
+      endpoint: "https://app.langwatch.ai",
+      maxEntries: 5,
+      negativeTtlMs: 60_000,
+      fetchImpl: fetchImpl as unknown as typeof fetch,
+    });
+
+    for (let i = 0; i < 20; i++) {
+      await verifier.verify(`sk-spray-${i}`);
+    }
+    expect(fetchImpl).toHaveBeenCalledTimes(20);
+
+    // The newest key is still cached, so it costs nothing to re-check.
+    await verifier.verify("sk-spray-19");
+    expect(fetchImpl).toHaveBeenCalledTimes(20);
+
+    // The oldest was evicted to keep the map bounded, so it has to be asked
+    // again even though its negative TTL has not passed.
+    await verifier.verify("sk-spray-0");
+    expect(fetchImpl).toHaveBeenCalledTimes(21);
+  });
+});