examples/MCP-tools Public
Limit tool calls per client with a Durable Object #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. [[rate-limits]] covers the bucket itself. | ||
| + | ||
| ## 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. | ||
Added.rebasely/kb/concepts/rate-limits.md+29−0
| @@ -0,0 +1,29 @@ | ||
| +--- | ||
| +id: rate-limits | ||
| +type: concept | ||
| +owner: justin | ||
| +component: rate-limiter | ||
| +describes: | ||
| + - src/rate-limiter/** | ||
| +reviewed: 2026-10-15 | ||
| +--- | ||
| +# Rate limits | ||
| + | ||
| +Each client gets a token bucket: 30 tool calls to start, one more every two seconds, never more than 30. [[0004-per-client-rate-limit]] has the reasoning. | ||
| + | ||
| +## Where the pieces are | ||
| + | ||
| +- `src/rate-limiter/bucket.ts`: the arithmetic. `take(bucket, now)` refills by the time elapsed, spends one token if there is one, and returns the new bucket and a decision: `{ ok: true, remaining }` or `{ ok: false, retryAfterSeconds }`. Pure, so `test/rate-limit.test.ts` drives it with a fake clock. | ||
| +- `src/rate-limiter/index.ts`: the `RateLimiter` Durable Object, bound as `RATE_LIMITER`. Its one RPC method, `take()`, loads the bucket from storage, runs the arithmetic, saves the result and returns the decision. | ||
| +- `src/worker.ts`: names the object by the caller's `cf-connecting-ip` and turns a refusal into the message the model sees. See [[http-transport]]. | ||
| + | ||
| +## Behavior worth knowing | ||
| + | ||
| +- A refused call does not spend anything, so a client that waits the time it was told always gets through. | ||
| +- A clock that steps backwards adds no tokens. | ||
| +- After each call the object sets an alarm for the moment its bucket would be full. If nothing else happens first, the alarm deletes the object's storage, and the next call starts from a fresh, full bucket, which is the same state. | ||
| +- Durable Objects handle one request at a time, so concurrent calls from one client cannot both spend the last token. | ||
| + | ||
| +## Changing the limit | ||
| + | ||
| +Edit `LIMIT` in `bucket.ts`. Buckets already stored keep their token count and refill toward the new capacity. Update the numbers in [[0004-per-client-rate-limit]] only by writing a new decision that supersedes it. | ||
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, [[rate-limits]]) 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 | ||
Added.rebasely/kb/decisions/0004-per-client-rate-limit.md+29−0
| @@ -0,0 +1,29 @@ | ||
| +--- | ||
| +id: 0004-per-client-rate-limit | ||
| +type: decision | ||
| +owner: justin | ||
| +reviewed: 2026-10-15 | ||
| +status: accepted | ||
| +--- | ||
| +# 0004 Per-client rate limit in a Durable Object | ||
| + | ||
| +## Context | ||
| + | ||
| +The tools are cheap, but not free, and a model stuck in a loop will call one hundreds of times a minute. The usage counters in KV cannot stop that: they are approximate and eventually consistent by design ([[0003-usage-counts-in-kv]]). A limit has to be exact at the moment it is checked, which needs a single writer per client. | ||
| + | ||
| +Workers has a rate limiting binding, but its counters are per location and approximate, and it cannot tell us how long a refused client should wait. | ||
| + | ||
| +## Decision | ||
| + | ||
| +- One `RateLimiter` Durable Object per client, named by `cf-connecting-ip`. It holds a token bucket: 30 tool calls in a burst, refilling one token every two seconds. | ||
| +- Only `tools/call` spends tokens, and only after the arguments pass the input schema. `initialize` and `tools/list` are free. | ||
| +- A refused call returns an `isError` tool result that says how many seconds to wait, in line with [[0002-tool-errors-are-results]]. Not an HTTP 429: most clients surface a 429 as a transport failure the model never reads, or retry it at once. | ||
| +- When the bucket is full again the object deletes its storage from an alarm, so idle clients cost nothing. | ||
| +- If the Durable Object call fails, the tool call is admitted and the error logged. | ||
| + | ||
| +## Consequences | ||
| + | ||
| +- Every tool call now makes one RPC to a Durable Object before it runs. The object lives near the client's first call, so that hop is usually short. | ||
| +- Clients behind one NAT or proxy share a bucket. Without auth there is nothing better to key on; if the server ever takes tokens, the key should become the token. | ||
| +- Failing open means a broken limiter lets a loop through. We prefer that to failing every call when the limiter has a bad minute. | ||
| +- The stateless transport ([[0001-stateless-streamable-http]]) is unchanged: the limiter keeps state per client, not per session. | ||
Changed.rebasely/kb/guides/connecting-a-client.md+2−1
| @@ -2,7 +2,7 @@ | ||
| id: connecting-a-client | ||
| type: guide | ||
| owner: justin | ||
| -reviewed: 2026-10-13 | ||
| +reviewed: 2026-10-15 | ||
| --- | ||
| # Connecting a client | ||
| @@ -55,4 +55,5 @@ | ||
| - **406 Not Acceptable**: the client did not send both `application/json` and `text/event-stream` in `Accept`. | ||
| - **415 Unsupported Media Type**: the body was not sent as `application/json`. | ||
| - **405 on GET**: the client is trying to open a standalone SSE stream or resume a session. Pick a Streamable HTTP client mode that does not need one; older SSE-only clients are not supported. | ||
| +- **"Rate limit reached for this client"**: the client made more than 30 tool calls in a short burst. Wait the number of seconds the message gives; one call every two seconds is always allowed. Clients behind the same IP share the limit. | ||
| - **Tools missing**: check `GET /` answers, then run `tools/list` with curl to see what the server sends. | ||
Changed.rebasely/kb/index.md+3−1
| @@ -1,11 +1,12 @@ | ||
| # mcp-tools knowledge base | ||
| -mcp-tools is a small MCP server on Cloudflare Workers. Three read-only tools (`convert_units`, `time_in_zone`, `text_stats`) answer from their input alone; the Worker serves them over Streamable HTTP at `/mcp` and counts calls in KV. | ||
| +mcp-tools is a small MCP server on Cloudflare Workers. Three read-only tools (`convert_units`, `time_in_zone`, `text_stats`) answer from their input alone; the Worker serves them over Streamable HTTP at `/mcp`, limits each client's tool calls with a Durable Object and counts calls in KV. | ||
| ## How it works | ||
| - [[tool-design]]: what every tool has (name, description for the model, schemas, structured output) and how failures reach the model. | ||
| - [[http-transport]]: how a POST to `/mcp` becomes an MCP response, and what the Worker does around it. | ||
| +- [[rate-limits]]: the per-client token bucket and the Durable Object that holds it. | ||
| - [[usage-counts]]: what `/stats` counts, where, and how far to trust it. | ||
| ## Why it is this way | ||
| @@ -13,6 +14,7 @@ | ||
| - [[0001-stateless-streamable-http]]: a fresh server per request, no sessions. | ||
| - [[0002-tool-errors-are-results]]: failures come back as tool results the model can read. | ||
| - [[0003-usage-counts-in-kv]]: approximate counters in KV rather than a Durable Object. | ||
| +- [[0004-per-client-rate-limit]]: an exact limit per client IP, refused calls answered as tool results. | ||
| ## How to | ||
ChangedREADME.md+6−2
| @@ -2,7 +2,7 @@ | ||
| A small [Model Context Protocol](https://modelcontextprotocol.io) server on Cloudflare Workers. It gives models three exact, read-only tools for things they tend to get wrong: unit arithmetic, the current time somewhere else, and counting words. | ||
| -The server speaks Streamable HTTP at `/mcp`, holds no sessions, and counts tool calls per day in a KV namespace. | ||
| +The server speaks Streamable HTTP at `/mcp`, holds no sessions, limits each client to a burst of 30 tool calls with a Durable Object, and counts tool calls per day in a KV namespace. | ||
| ## Tools | ||
| @@ -109,13 +109,14 @@ | ||
| npm run dev # wrangler dev on http://localhost:8787 | ||
| ``` | ||
| -`npx wrangler deploy` deploys it. The `USAGE` namespace in `wrangler.jsonc` has no id, so the first deploy creates it. | ||
| +`npx wrangler deploy` deploys it. The `USAGE` namespace in `wrangler.jsonc` has no id, so the first deploy creates it; the `RateLimiter` class is created by the `v1` migration. | ||
| ## Design notes | ||
| - **Stateless.** Every POST gets a fresh `McpServer` and a `WebStandardStreamableHTTPServerTransport` with no session id and JSON responses. Any isolate can answer any request; nothing needs a Durable Object. | ||
| - **Tools are pure.** Each tool in `src/tools/` is a name, a description written for the model, zod input and output schemas, and a function from input to output. `src/server.ts` is the only file that knows about MCP. | ||
| - **Errors are results.** A tool throws `ToolError` with advice ("Use an IANA name such as Europe/Paris"), and the server returns it with `isError: true` so the model can read it and try again. | ||
| +- **A limit per client.** Each caller IP gets a `RateLimiter` Durable Object with a token bucket: 30 calls, then one every two seconds. A refused call is an `isError` result that says how long to wait, so the model backs off instead of the client seeing a 429. | ||
| - **Counting stays off the request path.** Usage writes go through `ctx.waitUntil`, after the response. The counts are approximate, which is fine for seeing what gets used. | ||
| The longer versions are in the knowledge base under `.rebasely/kb/`, and the component map is `.rebasely/architecture/system.yaml`. | ||
| @@ -128,6 +129,9 @@ | ||
| worker.ts routes: /mcp, /stats, CORS | ||
| server.ts createServer(): registers the tools with McpServer | ||
| usage.ts per-day counters in KV | ||
| + rate-limiter/ | ||
| + index.ts RateLimiter Durable Object | ||
| + bucket.ts token bucket arithmetic | ||
| tools/ | ||
| convert-units.ts | ||
| time-in-zone.ts | ||
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"] }] | ||
| } | ||
View runtime
1 component added, 1 changed, 1 connection added
- Added
- Changed
- Removed
What runs on Cloudflare and the storage it touches, without the libraries.
The view runtime at the head of this pull request, with what it adds, changes and removes from the merge base fcaf000. Select a component for what changed in it.
rate-limiter
Durable ObjectClear
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.
Added by this pull request.
Connections: mcp-server → rate-limiter added.
Takes part in A client over its rate limit, An agent calling a tool.
Changes
- Changed: ~ view runtime: components +rate-limiter
- Changed: ~ mcp-server: description
- Added: + component rate-limiter (durable-object)
- Added: + edge mcp-server → rate-limiter (rpc)
Knowledge base
- The HTTP transport Changedconcepts/http-transport.md · concept
- Rate limits Addedconcepts/rate-limits.md · concept
- Tool design Changedconcepts/tool-design.md · concept
- 0004 Per-client rate limit in a Durable Object Addeddecisions/0004-per-client-rate-limit.md · decision
- Connecting a client Changedguides/connecting-a-client.md · guide
- mcp-tools knowledge base Changedindex.md · index