[vitest-pool-workers] fix: reset() now clears ratelimit binding state between tests (#14409)

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Author: matingathani

.changeset/vitest-pool-workers-reset-ratelimit.md

@@ -0,0 +1,6 @@
+---
+"miniflare": patch
+"@cloudflare/vitest-pool-workers": patch
+---
+
+`reset()` from `cloudflare:test` now resets ratelimit binding state between tests. Previously, `RATE_LIMITERS` bindings retained their in-memory bucket counts across test boundaries, causing later tests in the same file to see stale rate-limit exhaustion state.

fixtures/vitest-pool-workers-examples/reset/src/env.d.ts

@@ -1,16 +1,16 @@
 /* eslint-disable */
-// Generated by Wrangler by running `wrangler types ./reset/src/env.d.ts -c ./reset/wrangler.jsonc --no-include-runtime` (hash: 524c60174bc1732153e84c1b8699023e)
-interface __BaseEnv_Env {
-	KV_NAMESPACE: KVNamespace;
-	R2_BUCKET: R2Bucket;
-	DATABASE: D1Database;
-	COUNTER: DurableObjectNamespace<import("./index").Counter>;
-}
+// Generated by Wrangler by running `wrangler types src/env.d.ts -c wrangler.jsonc --no-include-runtime` (hash: 6992c0652c326145f228d37275b18131)
 declare namespace Cloudflare {
 	interface GlobalProps {
 		mainModule: typeof import("./index");
 		durableNamespaces: "Counter";
 	}
-	interface Env extends __BaseEnv_Env {}
+	interface Env {
+		KV_NAMESPACE: KVNamespace;
+		R2_BUCKET: R2Bucket;
+		DATABASE: D1Database;
+		RATE_LIMITER: RateLimit;
+		COUNTER: DurableObjectNamespace<import("./index").Counter>;
+	}
 }
-interface Env extends __BaseEnv_Env {}
+interface Env extends Cloudflare.Env {}

fixtures/vitest-pool-workers-examples/reset/test/reset.test.ts

@@ -62,3 +62,24 @@ it("sees reset Durable Object storage after reset", async ({ expect }) => {
 	const response = await stub.fetch("https://example.com");
 	expect(await response.text()).toBe("1");
 });
+
+it("exhausts ratelimit then resets between tests", async ({ expect }) => {
+	// First call succeeds
+	const first = await env.RATE_LIMITER.limit({ key: "test-key" });
+	expect(first.success).toBe(true);
+
+	// Exhaust the full limit of 100
+	for (let i = 1; i < 100; i++) {
+		await env.RATE_LIMITER.limit({ key: "test-key" });
+	}
+
+	// 101st call should fail
+	const over = await env.RATE_LIMITER.limit({ key: "test-key" });
+	expect(over.success).toBe(false);
+});
+
+it("sees reset ratelimit state after reset", async ({ expect }) => {
+	// After reset(), buckets should be cleared — first call succeeds again
+	const result = await env.RATE_LIMITER.limit({ key: "test-key" });
+	expect(result.success).toBe(true);
+});

fixtures/vitest-pool-workers-examples/reset/wrangler.jsonc

@@ -34,4 +34,11 @@
 			"database_id": "00000000-0000-0000-0000-000000000000",
 		},
 	],
+	"ratelimits": [
+		{
+			"name": "RATE_LIMITER",
+			"namespace_id": "1",
+			"simple": { "limit": 100, "period": 60 },
+		},
+	],
 }

packages/miniflare/src/plugins/ratelimit/index.ts

@@ -1,7 +1,20 @@
-import SCRIPT_RATELIMIT_OBJECT from "worker:ratelimit/ratelimit";
+import SCRIPT_RATELIMIT_CLIENT from "worker:ratelimit/ratelimit";
+import SCRIPT_RATELIMIT_OBJECT from "worker:ratelimit/ratelimit-object";
 import { z } from "zod";
-import { ProxyNodeBinding } from "../shared";
-import type { Worker_Binding } from "../../runtime";
+import { kVoid } from "../../runtime";
+import { SharedBindings } from "../../workers";
+import {
+	getMiniflareObjectBindings,
+	getUserBindingServiceName,
+	objectEntryWorker,
+	ProxyNodeBinding,
+	SERVICE_LOOPBACK,
+} from "../shared";
+import type {
+	Service,
+	Worker_Binding,
+	Worker_Binding_DurableObjectNamespaceDesignator,
+} from "../../runtime";
 import type { Plugin } from "../shared";
 
 export enum PeriodType {
@@ -14,7 +27,7 @@ export const RatelimitConfigSchema = z.object({
 		limit: z.number().gt(0),
 
 		// may relax this to be any number in the future
-		period: z.nativeEnum(PeriodType).optional(),
+		period: z.nativeEnum(PeriodType).optional().default(PeriodType.MINUTE),
 	}),
 });
 export const RatelimitOptionsSchema = z.object({
@@ -24,6 +37,12 @@ export const RatelimitOptionsSchema = z.object({
 export const RATELIMIT_PLUGIN_NAME = "ratelimit";
 const SERVICE_RATELIMIT_PREFIX = `${RATELIMIT_PLUGIN_NAME}`;
 const SERVICE_RATELIMIT_MODULE = `cloudflare-internal:${SERVICE_RATELIMIT_PREFIX}:module`;
+const RATELIMIT_ENTRY_SERVICE_PREFIX = `${RATELIMIT_PLUGIN_NAME}:ns`;
+const RATELIMIT_OBJECT_CLASS_NAME = "RateLimiterObject";
+const RATELIMIT_OBJECT: Worker_Binding_DurableObjectNamespaceDesignator = {
+	serviceName: RATELIMIT_ENTRY_SERVICE_PREFIX,
+	className: RATELIMIT_OBJECT_CLASS_NAME,
+};
 
 function buildJsonBindings(bindings: Record<string, any>): Worker_Binding[] {
 	return Object.entries(bindings).map(([name, value]) => ({
@@ -44,11 +63,21 @@ export const RATELIMIT_PLUGIN: Plugin<typeof RatelimitOptionsSchema> = {
 				name,
 				wrapped: {
 					moduleName: SERVICE_RATELIMIT_MODULE,
-					innerBindings: buildJsonBindings({
-						namespaceId: name,
-						limit: config.simple.limit,
-						period: config.simple.period,
-					}),
+					innerBindings: [
+						{
+							name: "fetcher",
+							service: {
+								name: getUserBindingServiceName(
+									RATELIMIT_ENTRY_SERVICE_PREFIX,
+									name
+								),
+							},
+						},
+						...buildJsonBindings({
+							limit: config.simple.limit,
+							period: config.simple.period,
+						}),
+					],
 				},
 			})
 		);
@@ -65,8 +94,46 @@ export const RATELIMIT_PLUGIN: Plugin<typeof RatelimitOptionsSchema> = {
 			])
 		);
 	},
-	async getServices() {
-		return [];
+	async getServices({ options, unsafeStickyBlobs }) {
+		if (!options.ratelimits) {
+			return [];
+		}
+		const names = Object.keys(options.ratelimits);
+		const services = names.map<Service>((name) => ({
+			name: getUserBindingServiceName(RATELIMIT_ENTRY_SERVICE_PREFIX, name),
+			worker: objectEntryWorker(RATELIMIT_OBJECT, name),
+		}));
+
+		const uniqueKey = `miniflare-${RATELIMIT_OBJECT_CLASS_NAME}`;
+		services.push({
+			name: RATELIMIT_ENTRY_SERVICE_PREFIX,
+			worker: {
+				compatibilityDate: "2023-07-24",
+				compatibilityFlags: ["nodejs_compat", "experimental"],
+				modules: [
+					{
+						name: "ratelimit-object.worker.js",
+						esModule: SCRIPT_RATELIMIT_OBJECT(),
+					},
+				],
+				durableObjectNamespaces: [
+					{ className: RATELIMIT_OBJECT_CLASS_NAME, uniqueKey },
+				],
+				// Counters are only ever needed for the lifetime of the Miniflare
+				// process, never persisted across restarts (matching the previous
+				// purely in-memory implementation).
+				durableObjectStorage: { inMemory: kVoid },
+				bindings: [
+					{
+						name: SharedBindings.MAYBE_SERVICE_LOOPBACK,
+						service: { name: SERVICE_LOOPBACK },
+					},
+					...getMiniflareObjectBindings(unsafeStickyBlobs),
+				],
+			},
+		});
+
+		return services;
 	},
 	getExtensions({ options }) {
 		if (!options.some((o) => o.ratelimits)) {
@@ -77,7 +144,7 @@ export const RATELIMIT_PLUGIN: Plugin<typeof RatelimitOptionsSchema> = {
 				modules: [
 					{
 						name: SERVICE_RATELIMIT_MODULE,
-						esModule: SCRIPT_RATELIMIT_OBJECT(),
+						esModule: SCRIPT_RATELIMIT_CLIENT(),
 						internal: true,
 					},
 				],

packages/miniflare/src/workers/ratelimit/ratelimit-object.worker.ts

@@ -0,0 +1,39 @@
+// Durable Object backing the emulated Ratelimit binding. Moving bucket/epoch
+// state here (instead of an in-memory object behind a `wrapped` binding)
+// means it's automatically torn down by `deleteAllDurableObjects()`, so
+// vitest-pool-workers' `reset()` clears it for free.
+import { MiniflareDurableObject, POST } from "miniflare:shared";
+import type { RouteHandler } from "miniflare:shared";
+
+interface LimitRequestBody {
+	key: string;
+	limit: number;
+	period: number;
+}
+
+interface LimitResult {
+	success: boolean;
+}
+
+export class RateLimiterObject extends MiniflareDurableObject {
+	#buckets = new Map<string, number>();
+	#epoch = 0;
+
+	@POST("/limit")
+	limit: RouteHandler = async (req) => {
+		const { key, limit, period } = await req.json<LimitRequestBody>();
+
+		const epoch = Math.floor(Date.now() / (period * 1000));
+		if (epoch !== this.#epoch) {
+			this.#epoch = epoch;
+			this.#buckets.clear();
+		}
+
+		const value = this.#buckets.get(key) ?? 0;
+		if (value >= limit) {
+			return Response.json({ success: false } satisfies LimitResult);
+		}
+		this.#buckets.set(key, value + 1);
+		return Response.json({ success: true } satisfies LimitResult);
+	};
+}

packages/miniflare/src/workers/ratelimit/ratelimit.worker.ts

@@ -1,10 +1,14 @@
 // Emulated Ratelimit Binding
+//
+// Thin client: keeps option validation local (so error messages surface
+// synchronously in the caller's isolate), forwards each `.limit()` call to
+// the `RateLimiterObject` Durable Object, which owns the bucket/epoch state.
 
 // ENV configuration
 interface RatelimitConfig {
-	namespaceId: number;
 	limit: number;
 	period: number;
+	fetcher: Fetcher;
 }
 
 // options for Ratelimit
@@ -25,22 +29,18 @@ function validate(test: boolean, message: string): asserts test {
 }
 
 class Ratelimit {
-	namespaceId: number;
+	fetcher: Fetcher;
 	limitVal: number;
 	period: number;
-	buckets: Map<string, number>;
-	epoch: number;
 
 	constructor(config: RatelimitConfig) {
-		this.namespaceId = config.namespaceId;
+		this.fetcher = config.fetcher;
 		this.limitVal = config.limit;
 		this.period = config.period;
-
-		this.buckets = new Map<string, number>();
-		this.epoch = 0;
 	}
 
-	// method that counts and checks against the limit in in-memory buckets
+	// method that counts and checks against the limit, delegating the actual
+	// bucket/epoch bookkeeping to the `RateLimiterObject` Durable Object
 	async limit(options: unknown): Promise<RatelimitResult> {
 		// validate options input
 		validate(
@@ -67,22 +67,11 @@ class Ratelimit {
 			`unsupported period: ${period}`
 		);
 
-		const epoch = Math.floor(Date.now() / (period * 1000));
-		if (epoch != this.epoch) {
-			// clear counters
-			this.epoch = epoch;
-			this.buckets.clear();
-		}
-		const val = this.buckets.get(key) || 0;
-		if (val >= limit) {
-			return {
-				success: false,
-			};
-		}
-		this.buckets.set(key, val + 1);
-		return {
-			success: true,
-		};
+		const res = await this.fetcher.fetch("http://ratelimit/limit", {
+			method: "POST",
+			body: JSON.stringify({ key, limit, period }),
+		});
+		return await res.json<RatelimitResult>();
 	}
 }