examples/MCP-tools Public
Limit tool calls per client with a Durable Object token bucket
A model stuck in a loop can call the same tool hundreds of times a minute, and nothing stopped it. Each client IP now gets a RateLimiter object with a bucket of 30 calls that refills one every two seconds. The server asks it before running a tool, after schema validation, and a refusal comes back as an isError result with the wait so the model can back off. If the limiter fails, the call goes through. The architecture gains the rate-limiter component and its RPC edge, the tool-call flow gains the token check, and a new flow shows the refused call.
Changed files
16 files changed266 additions17 deletionsagainst fcaf000
- Added.rebasely/architecture/flows/rate-limited-call.yaml+13−0
- Changed.rebasely/architecture/flows/tool-call.yaml+2−0
- Changed.rebasely/architecture/system.yaml+11−1
- Changed.rebasely/architecture/views/runtime.yaml+1−1
- Changed.rebasely/kb/concepts/http-transport.md+7−3
- Changed.rebasely/kb/concepts/tool-design.md+2−2
- Changedsrc/index.ts+3−1
- Addedsrc/rate-limiter/bucket.ts+25−0
- Addedsrc/rate-limiter/index.ts+20−0
- Changedsrc/server.ts+7−1
- Changedsrc/worker.ts+24−1
- Changedtest/fakes.ts+22−0
- Addedtest/rate-limit.test.ts+48−0
- Changedtest/server.test.ts+18−0
- Changedtest/worker.test.ts+58−6
- Changedwrangler.jsonc+5−1
Added.rebasely/architecture/flows/rate-limited-call.yaml+13−0
| @@ -0,0 +1,13 @@ | ||
| +title: A client over its rate limit | ||
| +description: >- | ||
| + An agent stuck in a loop has spent its 30 calls. The refusal comes back as a tool result, | ||
| + so the model reads how long to wait instead of the client reporting an HTTP error. | ||
| +actors: [agent, mcp-client] | ||
| +steps: | ||
| + - { from: agent, to: mcp-client, label: "call convert_units again, the 31st time this minute" } | ||
| + - { from: mcp-client, to: mcp-server, label: "POST /mcp: tools/call convert_units" } | ||
| + - { from: mcp-server, to: mcp-sdk, label: check the arguments against the input schema } | ||
| + - { from: mcp-server, to: rate-limiter, label: "take a token from this client's bucket" } | ||
| + - { from: rate-limiter, to: mcp-server, label: "refused: next token in 2 seconds" } | ||
| + - { from: mcp-server, to: mcp-client, label: "isError result: wait 2 seconds before the next call" } | ||
| + - { from: mcp-client, to: agent, label: "the refusal as a tool message, so the model can back off" } | ||
Changed.rebasely/architecture/flows/tool-call.yaml+2−0
| @@ -13,6 +13,8 @@ | ||
| - { from: agent, to: mcp-client, label: "call time_in_zone with zone tokyo" } | ||
| - { from: mcp-client, to: mcp-server, label: "POST /mcp: tools/call time_in_zone" } | ||
| - { from: mcp-server, to: mcp-sdk, label: check the arguments against the input schema } | ||
| + - { from: mcp-server, to: rate-limiter, label: "take a token from this client's bucket" } | ||
| + - { from: rate-limiter, to: mcp-server, label: "allowed, 29 tokens left" } | ||
| - { from: mcp-server, to: tools, label: "run it: resolve tokyo to Asia/Tokyo, read the offset" } | ||
| - { from: tools, to: mcp-server, label: "local time, or a ToolError saying what to fix" } | ||
| - { from: mcp-server, to: mcp-client, label: "structuredContent and JSON text, isError on failure" } | ||
Changed.rebasely/architecture/system.yaml+11−1
| @@ -3,7 +3,8 @@ | ||
| kind: worker | ||
| description: >- | ||
| The Worker. Serves MCP over Streamable HTTP at /mcp, statelessly: each POST gets a fresh | ||
| - McpServer with every tool registered. Also counts tool calls and serves them at /stats. | ||
| + McpServer with every tool registered. Charges each tool call to the caller's rate limit, | ||
| + counts tool calls and serves the counts at /stats. | ||
| owns: | ||
| - src/index.ts | ||
| - src/worker.ts | ||
| @@ -13,6 +14,15 @@ | ||
| - { to: tools, via: import } | ||
| - { to: mcp-sdk, via: import } | ||
| - { to: usage, via: kv } | ||
| + - { to: rate-limiter, via: rpc } | ||
| + rate-limiter: | ||
| + kind: durable-object | ||
| + description: >- | ||
| + One object per client IP holding a token bucket: 30 tool calls in a burst, one more every | ||
| + two seconds. Runs one call at a time, so the limit is exact. Clears itself once full. | ||
| + binding: RATE_LIMITER | ||
| + owns: | ||
| + - src/rate-limiter/** | ||
| tools: | ||
| kind: library | ||
| description: >- | ||
Changed.rebasely/architecture/views/runtime.yaml+1−1
| @@ -1,2 +1,2 @@ | ||
| description: What runs on Cloudflare and the storage it touches, without the libraries. | ||
| -components: [mcp-server, tools, usage] | ||
| +components: [mcp-server, tools, rate-limiter, usage] | ||
Changed.rebasely/kb/concepts/http-transport.md+7−3
| @@ -6,11 +6,11 @@ | ||
| describes: | ||
| - src/index.ts | ||
| - src/worker.ts | ||
| -reviewed: 2026-10-13 | ||
| +reviewed: 2026-10-15 | ||
| --- | ||
| # The HTTP transport | ||
| -`src/index.ts` is the Worker entry and only points `fetch` at `handle` in `src/worker.ts`. Keeping the handler out of the entry file means the tests import it under plain Node. | ||
| +`src/index.ts` is the Worker entry: it points `fetch` at `handle` in `src/worker.ts` and exports the `RateLimiter` Durable Object class. Keeping the handler out of the entry file means the tests import it under plain Node, without `cloudflare:workers`. | ||
| ## Routes | ||
| @@ -32,6 +32,10 @@ | ||
| Every `/mcp` and `/stats` response carries `access-control-allow-origin: *` and allows the MCP headers (`mcp-protocol-version`, `mcp-session-id`, `last-event-id`). The server has no cookies or credentials, so an open origin costs nothing and lets the MCP Inspector and web agents connect directly. | ||
| +## Rate limit | ||
| + | ||
| +`createServer` takes an `admit` hook, asked once per tool call after the arguments pass the input schema. The Worker answers it from the caller's `RateLimiter` object, named by the `cf-connecting-ip` header. A refusal comes back to the model as an `isError` result saying how many seconds to wait; `initialize` and `tools/list` are never charged. If the Durable Object call throws, the call is admitted and the error logged. | ||
| + | ||
| ## Usage counting | ||
| -`createServer` takes an `onCall` hook. The Worker hands it a function that passes the KV write to `ctx.waitUntil`, so the count happens after the response has gone and a KV error is logged rather than returned. | ||
| +`createServer` also takes an `onCall` hook, called only for calls that were admitted. The Worker hands it a function that passes the KV write to `ctx.waitUntil`, so the count happens after the response has gone and a KV error is logged rather than returned. | ||
Changed.rebasely/kb/concepts/tool-design.md+2−2
| @@ -6,7 +6,7 @@ | ||
| describes: | ||
| - src/tools/** | ||
| - src/server.ts | ||
| -reviewed: 2026-10-13 | ||
| +reviewed: 2026-10-15 | ||
| --- | ||
| # Tool design | ||
| @@ -39,7 +39,7 @@ | ||
| Unknown time zone "Mars/Olympus_Mons". Use an IANA name such as Europe/Paris, Asia/Tokyo or America/New_York. | ||
| ``` | ||
| -Arguments that fail the input schema never reach `run`; the SDK answers them with an `isError` result too. Any other exception is a bug: the server logs it and tells the model that retrying the same input will not help. | ||
| +Arguments that fail the input schema never reach `run`; the SDK answers them with an `isError` result too. A call the `admit` hook refuses (the per-client rate limit) is answered the same way, with how long to wait, and the tool never runs. Any other exception is a bug: the server logs it and tells the model that retrying the same input will not help. | ||
| ## Annotations | ||
Changedsrc/index.ts+3−1
| @@ -1,5 +1,7 @@ | ||
| // The Worker's entry point. The handler lives in worker.ts so tests can import it without | ||
| -// anything that only loads in workerd. | ||
| +// anything that only loads in workerd, such as the Durable Object class below. | ||
| import { handle, type Env } from "./worker.ts"; | ||
| + | ||
| +export { RateLimiter } from "./rate-limiter/index.ts"; | ||
| export default { fetch: handle } satisfies ExportedHandler<Env>; | ||
Addedsrc/rate-limiter/bucket.ts+25−0
| @@ -0,0 +1,25 @@ | ||
| +// A token bucket: each client starts with `capacity` tool calls, spends one per call and gets | ||
| +// them back at `refillPerSecond`. Pure, so the tests run it without a Durable Object. | ||
| + | ||
| +export type Limit = { capacity: number; refillPerSecond: number }; | ||
| +export type Bucket = { tokens: number; updated: number }; // updated: ms since the epoch | ||
| +export type Decision = { ok: true; remaining: number } | { ok: false; retryAfterSeconds: number }; | ||
| + | ||
| +// 30 calls in a burst, back to full after a minute of quiet. Plenty for an agent working | ||
| +// through a task; a loop calling the same tool forever runs dry quickly. | ||
| +export const LIMIT: Limit = { capacity: 30, refillPerSecond: 0.5 }; | ||
| + | ||
| +export function take(prev: Bucket | undefined, now: number, limit: Limit = LIMIT): { bucket: Bucket; decision: Decision } { | ||
| + const elapsed = prev ? Math.max(0, now - prev.updated) / 1000 : 0; | ||
| + const tokens = prev ? Math.min(limit.capacity, prev.tokens + elapsed * limit.refillPerSecond) : limit.capacity; | ||
| + if (tokens >= 1) { | ||
| + return { bucket: { tokens: tokens - 1, updated: now }, decision: { ok: true, remaining: Math.floor(tokens - 1) } }; | ||
| + } | ||
| + const retryAfterSeconds = Math.ceil((1 - tokens) / limit.refillPerSecond); | ||
| + return { bucket: { tokens, updated: now }, decision: { ok: false, retryAfterSeconds } }; | ||
| +} | ||
| + | ||
| +// How long until an untouched bucket is full again, at which point it is the same as no bucket. | ||
| +export function msUntilFull(bucket: Bucket, limit: Limit = LIMIT): number { | ||
| + return Math.ceil(((limit.capacity - bucket.tokens) / limit.refillPerSecond) * 1000); | ||
| +} | ||
Addedsrc/rate-limiter/index.ts+20−0
| @@ -0,0 +1,20 @@ | ||
| +import { DurableObject } from "cloudflare:workers"; | ||
| +import { msUntilFull, take, type Bucket, type Decision } from "./bucket.ts"; | ||
| + | ||
| +// One RateLimiter per client, named by the client's IP (src/worker.ts). A Durable Object runs | ||
| +// one call at a time, so two requests can never spend the same token. | ||
| +export class RateLimiter extends DurableObject { | ||
| + async take(): Promise<Decision> { | ||
| + const now = Date.now(); | ||
| + const { bucket, decision } = take(await this.ctx.storage.get<Bucket>("bucket"), now); | ||
| + await this.ctx.storage.put("bucket", bucket); | ||
| + // A full bucket is the same as an empty object, so clear it then instead of keeping one | ||
| + // row per IP we have ever seen. | ||
| + await this.ctx.storage.setAlarm(now + msUntilFull(bucket)); | ||
| + return decision; | ||
| + } | ||
| + | ||
| + async alarm(): Promise<void> { | ||
| + await this.ctx.storage.deleteAll(); | ||
| + } | ||
| +} | ||
Changedsrc/server.ts+7−1
| @@ -14,6 +14,9 @@ | ||
| export type ServerOptions = { | ||
| // The clock tools read; tests pin it. | ||
| now?: () => Date; | ||
| + // Asked before each tool call, after the arguments passed the input schema. A string refuses | ||
| + // the call and is returned to the model as an isError result; null lets it through. | ||
| + admit?: (tool: string) => Promise<string | null>; | ||
| // Called once per tool call that reached the tool, after it ran. `ok` is false when it | ||
| // returned an error result. | ||
| onCall?: (tool: string, ok: boolean) => void; | ||
| @@ -39,7 +42,10 @@ | ||
| outputSchema: tool.output, | ||
| annotations: { readOnlyHint: true, idempotentHint: true, openWorldHint: false }, | ||
| }, | ||
| - (args) => call(tool, args, options), | ||
| + async (args) => { | ||
| + const refusal = await options.admit?.(tool.name); | ||
| + return refusal ? failure(refusal) : call(tool, args, options); | ||
| + }, | ||
| ); | ||
| } | ||
Changedsrc/worker.ts+24−1
| @@ -1,8 +1,9 @@ | ||
| import { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js"; | ||
| +import type { RateLimiter } from "./rate-limiter/index.ts"; | ||
| import { createServer } from "./server.ts"; | ||
| import { MAX_DAYS, countCall, readUsage } from "./usage.ts"; | ||
| -export type Env = { USAGE: KVNamespace }; | ||
| +export type Env = { USAGE: KVNamespace; RATE_LIMITER: DurableObjectNamespace<RateLimiter> }; | ||
| // Browser clients (the MCP Inspector in direct mode, web agents) call /mcp cross-origin. The | ||
| // server has no cookies or credentials to protect, so any origin may call it. | ||
| @@ -36,12 +37,34 @@ | ||
| // A fresh server and transport for every request. Nothing carries over, so any isolate can | ||
| // answer any request and there is no session to lose. | ||
| // Counting happens after the response is sent; a KV hiccup never fails a tool call. | ||
| + const client = clientOf(request); | ||
| const server = createServer({ | ||
| + admit: () => admit(env, client), | ||
| onCall: (tool, ok) => ctx.waitUntil(countCall(env.USAGE, tool, ok).catch((err) => console.error("usage count failed", err))), | ||
| }); | ||
| const transport = new WebStandardStreamableHTTPServerTransport({ sessionIdGenerator: undefined, enableJsonResponse: true }); | ||
| await server.connect(transport); | ||
| return transport.handleRequest(request); | ||
| +} | ||
| + | ||
| +// Who is calling, for the rate limit. Cloudflare sets cf-connecting-ip on every request that | ||
| +// reaches the Worker; without it (local tests, odd proxies) all such callers share one bucket. | ||
| +function clientOf(request: Request): string { | ||
| + return request.headers.get("cf-connecting-ip") ?? "unknown"; | ||
| +} | ||
| + | ||
| +// Spends one of the client's tokens. If the limiter itself fails, the call goes through: a | ||
| +// broken limit should cost us a little abuse, not every caller their answers. | ||
| +async function admit(env: Env, client: string): Promise<string | null> { | ||
| + try { | ||
| + const decision = await env.RATE_LIMITER.getByName(client).take(); | ||
| + if (decision.ok) return null; | ||
| + const wait = decision.retryAfterSeconds; | ||
| + return `Rate limit reached for this client. Wait ${wait} second${wait === 1 ? "" : "s"} before the next tool call; calling sooner will be refused the same way.`; | ||
| + } catch (err) { | ||
| + console.error("rate limiter failed; admitting the call", err); | ||
| + return null; | ||
| + } | ||
| } | ||
| // GET /stats?days=7: calls and errors per tool for the last few UTC days, newest first. | ||
Changedtest/fakes.ts+22−0
| @@ -1,5 +1,8 @@ | ||
| // Small stand-ins for the Workers runtime, enough for the code under test. | ||
| +import { take, type Bucket } from "../src/rate-limiter/bucket.ts"; | ||
| +import type { RateLimiter } from "../src/rate-limiter/index.ts"; | ||
| + | ||
| type Entry = { value: string; metadata: unknown }; | ||
| // The parts of KVNamespace that src/usage.ts uses, over a Map. list() pages by `limit` so the | ||
| @@ -41,4 +44,23 @@ | ||
| settle: async () => void (await Promise.all(pending.splice(0))), | ||
| }; | ||
| return ctx as unknown as ExecutionContext & { settle: () => Promise<void> }; | ||
| +} | ||
| + | ||
| +// A RATE_LIMITER namespace with one in-memory bucket per name, running the real bucket logic | ||
| +// against a clock the test controls. | ||
| +export function fakeLimiter(clock = { now: 0 }): DurableObjectNamespace<RateLimiter> & { clock: { now: number } } { | ||
| + const buckets = new Map<string, Bucket>(); | ||
| + const ns = { | ||
| + clock, | ||
| + getByName(name: string) { | ||
| + return { | ||
| + async take() { | ||
| + const { bucket, decision } = take(buckets.get(name), clock.now); | ||
| + buckets.set(name, bucket); | ||
| + return decision; | ||
| + }, | ||
| + }; | ||
| + }, | ||
| + }; | ||
| + return ns as unknown as DurableObjectNamespace<RateLimiter> & { clock: { now: number } }; | ||
| } | ||
Addedtest/rate-limit.test.ts+48−0
| @@ -0,0 +1,48 @@ | ||
| +import assert from "node:assert/strict"; | ||
| +import { test } from "node:test"; | ||
| +import { LIMIT, msUntilFull, take, type Bucket } from "../src/rate-limiter/bucket.ts"; | ||
| + | ||
| +const limit = { capacity: 3, refillPerSecond: 1 }; | ||
| + | ||
| +function drain(n: number, now = 0): Bucket | undefined { | ||
| + let bucket: Bucket | undefined; | ||
| + for (let i = 0; i < n; i++) bucket = take(bucket, now, limit).bucket; | ||
| + return bucket; | ||
| +} | ||
| + | ||
| +test("a new client starts with a full bucket", () => { | ||
| + assert.deepEqual(take(undefined, 0, limit).decision, { ok: true, remaining: 2 }); | ||
| +}); | ||
| + | ||
| +test("an empty bucket refuses and says how long to wait", () => { | ||
| + const { bucket, decision } = take(drain(3), 0, limit); | ||
| + assert.deepEqual(decision, { ok: false, retryAfterSeconds: 1 }); | ||
| + assert.equal(bucket.tokens, 0, "a refused call costs nothing"); | ||
| +}); | ||
| + | ||
| +test("tokens come back over time, up to the capacity", () => { | ||
| + const empty = drain(3); | ||
| + assert.deepEqual(take(empty, 1000, limit).decision, { ok: true, remaining: 0 }); | ||
| + assert.deepEqual(take(empty, 60_000, limit).decision, { ok: true, remaining: 2 }); | ||
| +}); | ||
| + | ||
| +test("a partly refilled bucket rounds the wait up", () => { | ||
| + const empty = drain(3); | ||
| + const slow = { capacity: 3, refillPerSecond: 0.5 }; | ||
| + assert.deepEqual(take(empty, 500, slow).decision, { ok: false, retryAfterSeconds: 2 }); | ||
| +}); | ||
| + | ||
| +test("a clock that goes backwards does not mint tokens", () => { | ||
| + const bucket = drain(3, 10_000); | ||
| + assert.equal(take(bucket, 5_000, limit).decision.ok, false); | ||
| +}); | ||
| + | ||
| +test("knows when a bucket is full again", () => { | ||
| + assert.equal(msUntilFull({ tokens: 1, updated: 0 }, limit), 2000); | ||
| + assert.equal(msUntilFull({ tokens: 0, updated: 0 }), 60_000); | ||
| +}); | ||
| + | ||
| +test("the default allows a burst of 30 and one call every two seconds after", () => { | ||
| + assert.equal(LIMIT.capacity, 30); | ||
| + assert.equal(1 / LIMIT.refillPerSecond, 2); | ||
| +}); | ||
Changedtest/server.test.ts+18−0
| @@ -58,6 +58,24 @@ | ||
| await client.close(); | ||
| }); | ||
| +test("admit can refuse a call before the tool runs", async () => { | ||
| + const asked: string[] = []; | ||
| + const calls: string[] = []; | ||
| + const client = await connect({ | ||
| + admit: async (tool) => (asked.push(tool), tool === "text_stats" ? "Rate limit reached for this client. Wait 4 seconds." : null), | ||
| + onCall: (tool) => calls.push(tool), | ||
| + }); | ||
| + const refused = await client.callTool({ name: "text_stats", arguments: { text: "Hi." } }); | ||
| + assert.equal(refused.isError, true); | ||
| + assert.equal(textOf(refused), "Rate limit reached for this client. Wait 4 seconds."); | ||
| + const allowed = await client.callTool({ name: "convert_units", arguments: { value: 1, from: "m", to: "cm" } }); | ||
| + assert.equal(allowed.isError, undefined); | ||
| + await client.callTool({ name: "convert_units", arguments: { value: "bad", from: "m", to: "cm" } }); | ||
| + assert.deepEqual(asked, ["text_stats", "convert_units"], "arguments that fail the schema are never charged"); | ||
| + assert.deepEqual(calls, ["convert_units"]); | ||
| + await client.close(); | ||
| +}); | ||
| + | ||
| test("a ToolError comes back as an isError result, not a protocol error", async () => { | ||
| const client = await connect(); | ||
| const result = await client.callTool({ name: "convert_units", arguments: { value: 1, from: "kg", to: "m" } }); | ||
Changedtest/worker.test.ts+58−6
| @@ -1,15 +1,20 @@ | ||
| import assert from "node:assert/strict"; | ||
| import { test } from "node:test"; | ||
| import { handle, type Env } from "../src/worker.ts"; | ||
| -import { fakeCtx, fakeKv } from "./fakes.ts"; | ||
| +import { fakeCtx, fakeKv, fakeLimiter } from "./fakes.ts"; | ||
| -const env: Env = { USAGE: fakeKv() }; | ||
| +const env: Env = { USAGE: fakeKv(), RATE_LIMITER: fakeLimiter() }; | ||
| const ctx = fakeCtx(); | ||
| -function rpc(method: string, params?: unknown, id = 1): Request { | ||
| +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" }, | ||
| + 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 } : {}) }), | ||
| }); | ||
| } | ||
| @@ -47,7 +52,7 @@ | ||
| }); | ||
| test("tool calls are counted after the response, and /stats reports them", async () => { | ||
| - const env: Env = { USAGE: fakeKv() }; | ||
| + 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); | ||
| @@ -65,13 +70,60 @@ | ||
| 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 }, ctx); | ||
| + 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; | ||
| } | ||
Changedwrangler.jsonc+5−1
| @@ -6,5 +6,9 @@ | ||
| "main": "src/index.ts", | ||
| "compatibility_date": "2026-09-01", | ||
| "observability": { "enabled": true }, | ||
| - "kv_namespaces": [{ "binding": "USAGE" }] | ||
| + "kv_namespaces": [{ "binding": "USAGE" }], | ||
| + "durable_objects": { | ||
| + "bindings": [{ "name": "RATE_LIMITER", "class_name": "RateLimiter" }] | ||
| + }, | ||
| + "migrations": [{ "tag": "v1", "new_sqlite_classes": ["RateLimiter"] }] | ||
| } | ||
Checks
No checks reported for this commit.
Runs
No runs for this commit.