Fix workers AI usage with @cloudflare/vite-plugin

packages/vite-plugin-cloudflare/playground/ai/__tests__/basic-functionality.spec.ts

@@ -0,0 +1,9 @@
+import { expect, test } from "vitest";
+import { getTextResponse } from "../../__test-utils__";
+
+// todo(justinvdm, 2025-02-04): Since this example uses Workers AI, there is CLI interactivity required to choose and authenticate for the account
+// that should be used when talking to the Workers AI API, making it difficult to run this test automatically. Once the plugin accepts an account id
+// this test might be able to be unskipped.
+test.skip("basic hello-world functionality", async () => {
+	expect(JSON.parse(await getTextResponse())).toContain("PONG");
+});

packages/vite-plugin-cloudflare/playground/ai/package.json

@@ -0,0 +1,19 @@
+{
+	"name": "@playground/ai",
+	"private": true,
+	"type": "module",
+	"scripts": {
+		"build": "vite build --app",
+		"check:types": "tsc --build",
+		"dev": "vite dev",
+		"preview": "vite preview"
+	},
+	"devDependencies": {
+		"@cloudflare/vite-plugin": "workspace:*",
+		"@cloudflare/workers-tsconfig": "workspace:*",
+		"@cloudflare/workers-types": "^4.20250129.0",
+		"typescript": "catalog:default",
+		"vite": "catalog:vite-plugin",
+		"wrangler": "catalog:vite-plugin"
+	}
+}

packages/vite-plugin-cloudflare/playground/ai/src/index.ts

@@ -0,0 +1,15 @@
+export interface Env {
+	AI: Ai;
+}
+
+export default {
+	async fetch(_request: Request, env: Env): Promise<Response> {
+		const response = await env.AI.run("@cf/meta/llama-3.1-8b-instruct", {
+			prompt: "When I say PING, you say PONG. PING",
+		});
+
+		return new Response(
+			JSON.stringify((response as { response: string }).response.toUpperCase())
+		);
+	},
+} satisfies ExportedHandler<Env>;

packages/vite-plugin-cloudflare/playground/ai/tsconfig.json

@@ -0,0 +1,7 @@
+{
+	"files": [],
+	"references": [
+		{ "path": "./tsconfig.node.json" },
+		{ "path": "./tsconfig.worker.json" }
+	]
+}

packages/vite-plugin-cloudflare/playground/ai/tsconfig.node.json

@@ -0,0 +1,4 @@
+{
+	"extends": ["@cloudflare/workers-tsconfig/base.json"],
+	"include": ["vite.config.ts", "__tests__"]
+}

packages/vite-plugin-cloudflare/playground/ai/tsconfig.worker.json

@@ -0,0 +1,4 @@
+{
+	"extends": ["@cloudflare/workers-tsconfig/worker.json"],
+	"include": ["src"]
+}

packages/vite-plugin-cloudflare/playground/ai/turbo.json

@@ -0,0 +1,9 @@
+{
+	"$schema": "http://turbo.build/schema.json",
+	"extends": ["//"],
+	"tasks": {
+		"build": {
+			"outputs": ["dist/**"]
+		}
+	}
+}

packages/vite-plugin-cloudflare/playground/ai/vite.config.ts

@@ -0,0 +1,6 @@
+import { cloudflare } from "@cloudflare/vite-plugin";
+import { defineConfig } from "vite";
+
+export default defineConfig({
+	plugins: [cloudflare({ persistState: false })],
+});

packages/vite-plugin-cloudflare/playground/ai/wrangler.toml

@@ -0,0 +1,6 @@
+name = "worker"
+main = "./src/index.ts"
+compatibility_date = "2024-12-30"
+
+[ai]
+binding = "AI"

packages/vite-plugin-cloudflare/src/miniflare-options.ts

@@ -276,7 +276,7 @@ export function getDevMiniflareOptions(
 		},
 	];
 
-	const userWorkers =
+	const workersFromConfig =
 		resolvedPluginConfig.type === "workers"
 			? Object.entries(resolvedPluginConfig.workers).map(
 					([environmentName, workerConfig]) => {
@@ -288,69 +288,99 @@ export function getDevMiniflareOptions(
 							resolvedPluginConfig.cloudflareEnv
 						);
 
+						const { externalWorkers } = miniflareWorkerOptions;
+
 						const { ratelimits, ...workerOptions } =
 							miniflareWorkerOptions.workerOptions;
 
 						return {
-							...workerOptions,
-							// We have to add the name again because `unstable_getMiniflareWorkerOptions` sets it to `undefined`
-							name: workerConfig.name,
-							modulesRoot: miniflareModulesRoot,
-							unsafeEvalBinding: "__VITE_UNSAFE_EVAL__",
-							bindings: {
-								...workerOptions.bindings,
-								__VITE_ROOT__: resolvedViteConfig.root,
-								__VITE_ENTRY_PATH__: workerConfig.main,
-							},
-							serviceBindings: {
-								...workerOptions.serviceBindings,
-								...(environmentName ===
-									resolvedPluginConfig.entryWorkerEnvironmentName &&
-								workerConfig.assets?.binding
-									? {
-											[workerConfig.assets.binding]: ASSET_WORKER_NAME,
+							externalWorkers,
+							worker: {
+								...workerOptions,
+								// We have to add the name again because `unstable_getMiniflareWorkerOptions` sets it to `undefined`
+								name: workerConfig.name,
+								modulesRoot: miniflareModulesRoot,
+								unsafeEvalBinding: "__VITE_UNSAFE_EVAL__",
+								bindings: {
+									...workerOptions.bindings,
+									__VITE_ROOT__: resolvedViteConfig.root,
+									__VITE_ENTRY_PATH__: workerConfig.main,
+								},
+								serviceBindings: {
+									...workerOptions.serviceBindings,
+									...(environmentName ===
+										resolvedPluginConfig.entryWorkerEnvironmentName &&
+									workerConfig.assets?.binding
+										? {
+												[workerConfig.assets.binding]: ASSET_WORKER_NAME,
+											}
+										: {}),
+									__VITE_INVOKE_MODULE__: async (request) => {
+										const payload =
+											(await request.json()) as vite.CustomPayload;
+										const invokePayloadData = payload.data as {
+											id: string;
+											name: string;
+											data: [string, string, FetchFunctionOptions];
+										};
+
+										assert(
+											invokePayloadData.name === "fetchModule",
+											`Invalid invoke event: ${invokePayloadData.name}`
+										);
+
+										const [moduleId] = invokePayloadData.data;
+
+										// For some reason we need this here for cloudflare built-ins (e.g. `cloudflare:workers`) but not for node built-ins (e.g. `node:path`)
+										// See https://github.com/flarelabs-net/vite-plugin-cloudflare/issues/46
+										if (moduleId.startsWith("cloudflare:")) {
+											const result = {
+												externalize: moduleId,
+												type: "builtin",
+											} satisfies vite.FetchResult;
+
+											return new MiniflareResponse(JSON.stringify({ result }));
 										}
-									: {}),
-								__VITE_INVOKE_MODULE__: async (request) => {
-									const payload = (await request.json()) as vite.CustomPayload;
-									const invokePayloadData = payload.data as {
-										id: string;
-										name: string;
-										data: [string, string, FetchFunctionOptions];
-									};
-
-									assert(
-										invokePayloadData.name === "fetchModule",
-										`Invalid invoke event: ${invokePayloadData.name}`
-									);
-
-									const [moduleId] = invokePayloadData.data;
-
-									// For some reason we need this here for cloudflare built-ins (e.g. `cloudflare:workers`) but not for node built-ins (e.g. `node:path`)
-									// See https://github.com/flarelabs-net/vite-plugin-cloudflare/issues/46
-									if (moduleId.startsWith("cloudflare:")) {
-										const result = {
-											externalize: moduleId,
-											type: "builtin",
-										} satisfies vite.FetchResult;
-
-										return new MiniflareResponse(JSON.stringify({ result }));
-									}
-
-									const devEnvironment = viteDevServer.environments[
-										environmentName
-									] as CloudflareDevEnvironment;
-
-									const result = await devEnvironment.hot.handleInvoke(payload);
-
-									return new MiniflareResponse(JSON.stringify(result));
+
+										const devEnvironment = viteDevServer.environments[
+											environmentName
+										] as CloudflareDevEnvironment;
+
+										const result =
+											await devEnvironment.hot.handleInvoke(payload);
+
+										return new MiniflareResponse(JSON.stringify(result));
+									},
 								},
-							},
-						} satisfies Partial<WorkerOptions>;
+							} satisfies Partial<WorkerOptions>,
+						};
 					}
 				)
 			: [];
 
+	const userWorkers = workersFromConfig.map((worker) => worker.worker);
+
+	const resolvedExternalWorkersMap = new Map(
+		workersFromConfig
+			.flatMap((worker) => worker.externalWorkers)
+			.map((worker) => [worker.name, worker])
+	);
+
+	const externalWorkersMap = new Map<string, WorkerOptions>();
+
+	for (const worker of userWorkers) {
+		for (const binding of Object.values(worker.wrappedBindings ?? {})) {
+			const scriptName =
+				typeof binding === "string" ? binding : binding.scriptName;
+
+			const externalWorker = resolvedExternalWorkersMap.get(scriptName);
+
+			if (externalWorker) {
+				externalWorkersMap.set(scriptName, externalWorker);
+			}
+		}
+	}
+
 	const workerToWorkerEntrypointNamesMap =
 		getWorkerToWorkerEntrypointNamesMap(userWorkers);
 	const workerToDurableObjectClassNamesMap =
@@ -375,6 +405,7 @@ export function getDevMiniflareOptions(
 		),
 		workers: [
 			...assetWorkers,
+			...externalWorkersMap.values(),
 			...userWorkers.map((workerOptions) => {
 				const wrappers = [
 					`import { createWorkerEntrypointWrapper, createDurableObjectWrapper, createWorkflowEntrypointWrapper } from '${RUNNER_PATH}';`,

packages/wrangler/src/api/integrations/platform/index.ts

@@ -236,6 +236,7 @@ export interface Unstable_MiniflareWorkerOptions {
 	workerOptions: SourcelessWorkerOptions;
 	define: Record<string, string>;
 	main?: string;
+	externalWorkers: WorkerOptions[];
 }
 
 export function unstable_getMiniflareWorkerOptions(
@@ -264,7 +265,7 @@ export function unstable_getMiniflareWorkerOptions(
 		}));
 
 	const bindings = getBindings(config, env, true, {});
-	const { bindingOptions } = buildMiniflareBindingOptions({
+	const { bindingOptions, externalWorkers } = buildMiniflareBindingOptions({
 		name: undefined,
 		bindings,
 		workerDefinitions: undefined,
@@ -330,5 +331,10 @@ export function unstable_getMiniflareWorkerOptions(
 		...assetOptions,
 	};
 
-	return { workerOptions, define: config.define, main: config.main };
+	return {
+		workerOptions,
+		define: config.define,
+		main: config.main,
+		externalWorkers,
+	};
 }