Rebasely Beta: data may be purged

examples/MCP-tools Public

TypeScript · 142 lines · 6.4 KiB · at 2090b14

import assert from "node:assert/strict";import { test } from "node:test";import { handle, type Env } from "../src/worker.ts";import { fakeCtx, fakeKv, fakeLimiter } from "./fakes.ts";const env: Env = { USAGE: fakeKv(), RATE_LIMITER: fakeLimiter() };const ctx = fakeCtx();function rpc(method: string, params?: unknown, id = 1, ip = "203.0.113.7"): Request {  return new Request("http://localhost/mcp", {    method: "POST",    headers: {      "content-type": "application/json",      accept: "application/json, text/event-stream",      "mcp-protocol-version": "2025-06-18",      "cf-connecting-ip": ip,    },    body: JSON.stringify({ jsonrpc: "2.0", id, method, ...(params ? { params } : {}) }),  });}test("POST /mcp answers initialize as JSON", async () => {  const res = await handle(rpc("initialize", { protocolVersion: "2025-06-18", capabilities: {}, clientInfo: { name: "test", version: "0" } }), env, ctx);  assert.equal(res.status, 200);  assert.match(res.headers.get("content-type") ?? "", /application\/json/);  assert.equal(res.headers.get("mcp-session-id"), null);  const body = (await res.json()) as { result: { serverInfo: { name: string } } };  assert.equal(body.result.serverInfo.name, "mcp-tools");});test("each request gets its own server, so tools/call needs no session", async () => {  const res = await handle(rpc("tools/call", { name: "convert_units", arguments: { value: 1, from: "ft", to: "in" } }), env, ctx);  const body = (await res.json()) as { result: { structuredContent: { result: number } } };  assert.equal(body.result.structuredContent.result, 12);});test("GET and DELETE on /mcp are refused with 405", async () => {  for (const method of ["GET", "DELETE"]) {    const res = await handle(new Request("http://localhost/mcp", { method }), env, ctx);    assert.equal(res.status, 405);    assert.equal(res.headers.get("allow"), "POST, OPTIONS");  }});test("CORS preflights succeed and responses carry the CORS headers", async () => {  const pre = await handle(new Request("http://localhost/mcp", { method: "OPTIONS" }), env, ctx);  assert.equal(pre.status, 204);  assert.equal(pre.headers.get("access-control-allow-origin"), "*");  assert.match(pre.headers.get("access-control-allow-headers") ?? "", /mcp-protocol-version/);  const res = await handle(rpc("tools/list"), env, ctx);  assert.equal(res.headers.get("access-control-allow-origin"), "*");});test("tool calls are counted after the response, and /stats reports them", async () => {  const env: Env = { USAGE: fakeKv(), RATE_LIMITER: fakeLimiter() };  const ctx = fakeCtx();  await handle(rpc("tools/call", { name: "text_stats", arguments: { text: "Hello there." } }), env, ctx);  await handle(rpc("tools/call", { name: "time_in_zone", arguments: { zone: "Atlantis/Capital" } }), env, ctx);  await handle(rpc("tools/list"), env, ctx);  await ctx.settle();  const res = await handle(new Request("http://localhost/stats?days=1"), env, ctx);  assert.equal(res.status, 200);  const { days } = (await res.json()) as { days: { tools: Record<string, unknown> }[] };  assert.deepEqual(days[0].tools, { text_stats: { calls: 1, errors: 0 }, time_in_zone: { calls: 1, errors: 1 } });});test("a failing usage write does not fail the tool call", async () => {  const broken = fakeKv();  broken.put = async () => {    throw new Error("KV is down");  };  const ctx = fakeCtx();  const res = await handle(rpc("tools/call", { name: "convert_units", arguments: { value: 1, from: "m", to: "cm" } }), { USAGE: broken, RATE_LIMITER: fakeLimiter() }, ctx);  const body = (await res.json()) as { result: { isError?: boolean } };  assert.equal(body.result.isError, undefined);  const logged = console.error;  console.error = () => {};  try {    await ctx.settle();  } finally {    console.error = logged;  }});test("a client past its limit gets an isError result with the wait, and others are unaffected", async () => {  const limiter = fakeLimiter();  const env: Env = { USAGE: fakeKv(), RATE_LIMITER: limiter };  const ctx = fakeCtx();  const call = async (ip: string) => {    const res = await handle(rpc("tools/call", { name: "convert_units", arguments: { value: 1, from: "m", to: "mm" } }, 1, ip), env, ctx);    return ((await res.json()) as { result: { isError?: boolean; content: { text: string }[] } }).result;  };  for (let i = 0; i < 30; i++) assert.equal((await call("198.51.100.1")).isError, undefined);  const refused = await call("198.51.100.1");  assert.equal(refused.isError, true);  assert.match(refused.content[0].text, /Rate limit reached for this client\. Wait 2 seconds/);  assert.equal((await call("198.51.100.2")).isError, undefined);  limiter.clock.now += 2000;  assert.equal((await call("198.51.100.1")).isError, undefined);  await ctx.settle();  const res = await handle(new Request("http://localhost/stats?days=1"), env, ctx);  const { days } = (await res.json()) as { days: { tools: Record<string, { calls: number }> }[] };  assert.equal(days[0].tools.convert_units.calls, 32, "refused calls are not counted as usage");});test("only tool calls spend tokens", async () => {  const env: Env = { USAGE: fakeKv(), RATE_LIMITER: fakeLimiter() };  for (let i = 0; i < 40; i++) await handle(rpc("tools/list"), env, ctx);  const res = await handle(rpc("tools/call", { name: "text_stats", arguments: { text: "Still allowed." } }), env, ctx);  const body = (await res.json()) as { result: { isError?: boolean } };  assert.equal(body.result.isError, undefined);});test("a failing rate limiter lets the call through", async () => {  const limiter = fakeLimiter();  limiter.getByName = () => {    throw new Error("Durable Object overloaded");  };  const logged = console.error;  console.error = () => {};  try {    const res = await handle(rpc("tools/call", { name: "convert_units", arguments: { value: 2, from: "kg", to: "g" } }), { USAGE: fakeKv(), RATE_LIMITER: limiter }, fakeCtx());    const body = (await res.json()) as { result: { structuredContent: { result: number } } };    assert.equal(body.result.structuredContent.result, 2000);  } finally {    console.error = logged;  }});test("/stats checks its days parameter", async () => {  for (const days of ["0", "32", "two"]) {    const res = await handle(new Request(`http://localhost/stats?days=${days}`), env, ctx);    assert.equal(res.status, 400);  }});test("other paths are 404", async () => {  const res = await handle(new Request("http://localhost/nope"), env, ctx);  assert.equal(res.status, 404);});