Add support for Workers AI in local mode

.changeset/mean-jars-count.md

@@ -0,0 +1,5 @@
+---
+"wrangler": minor
+---
+
+Add support for Workers AI in local mode

fixtures/ai-app/.gitignore

@@ -0,0 +1 @@
+dist

fixtures/ai-app/package.json

@@ -0,0 +1,21 @@
+{
+	"name": "ai-app",
+	"version": "1.0.1",
+	"private": true,
+	"description": "",
+	"license": "ISC",
+	"author": "",
+	"main": "src/index.js",
+	"scripts": {
+		"check:type": "tsc",
+		"test": "vitest run",
+		"test:watch": "vitest",
+		"type:tests": "tsc -p ./tests/tsconfig.json"
+	},
+	"devDependencies": {
+		"undici": "^5.23.0",
+		"wrangler": "workspace:*",
+		"@cloudflare/workers-tsconfig": "workspace:^",
+		"@cloudflare/ai": "^1.0.35"
+	}
+}

fixtures/ai-app/src/index.js

@@ -0,0 +1,12 @@
+console.log("startup log");
+
+export default {
+	async fetch(request, env) {
+		console.log("request log");
+
+		return Response.json({
+			binding: env.AI,
+			fetcher: env.AI.fetch.toString(),
+		});
+	},
+};

fixtures/ai-app/tests/index.test.ts

@@ -0,0 +1,31 @@
+import { resolve } from "path";
+import { fetch } from "undici";
+import { describe, it, beforeAll, afterAll } from "vitest";
+import { runWranglerDev } from "../../shared/src/run-wrangler-long-lived";
+
+describe("'wrangler dev' correctly renders pages", () => {
+	let ip: string,
+		port: number,
+		stop: (() => Promise<unknown>) | undefined,
+		getOutput: () => string;
+
+	beforeAll(async () => {
+		({ ip, port, stop, getOutput } = await runWranglerDev(
+			resolve(__dirname, ".."),
+			["--local", "--port=0", "--inspector-port=0"]
+		));
+	});
+
+	afterAll(async () => {
+		await stop?.();
+	});
+
+	it("ai binding is defined ", async ({ expect }) => {
+		const response = await fetch(`http://${ip}:${port}/`);
+		const content = await response.json();
+		expect(content).toEqual({
+			binding: {},
+			fetcher: "function fetch() { [native code] }",
+		});
+	});
+});

fixtures/ai-app/tests/tsconfig.json

@@ -0,0 +1,7 @@
+{
+	"extends": "@cloudflare/workers-tsconfig/tsconfig.json",
+	"compilerOptions": {
+		"types": ["node"]
+	},
+	"include": ["**/*.ts", "../../../node-types.d.ts"]
+}

fixtures/ai-app/tsconfig.json

@@ -0,0 +1,13 @@
+{
+	"compilerOptions": {
+		"target": "ES2020",
+		"esModuleInterop": true,
+		"module": "CommonJS",
+		"lib": ["ES2020"],
+		"types": ["node"],
+		"skipLibCheck": true,
+		"moduleResolution": "node",
+		"noEmit": true
+	},
+	"include": ["tests", "../../node-types.d.ts"]
+}

fixtures/ai-app/vitest.config.ts

@@ -0,0 +1,10 @@
+import { defineConfig } from "vitest/config";
+
+export default defineConfig({
+	test: {
+		testTimeout: 10_000,
+		hookTimeout: 10_000,
+		teardownTimeout: 10_000,
+		useAtomics: true,
+	},
+});

fixtures/ai-app/wrangler.toml

@@ -0,0 +1,7 @@
+name = "ai-app"
+compatibility_date = "2023-11-21"
+
+main = "src/index.js"
+
+[ai]
+binding = "AI"

packages/wrangler/package.json

@@ -117,6 +117,7 @@
 		"xxhash-wasm": "^1.0.1"
 	},
 	"devDependencies": {
+		"@cloudflare/ai": "^1.0.35",
 		"@cloudflare/cli": "workspace:*",
 		"@cloudflare/eslint-config-worker": "*",
 		"@cloudflare/pages-shared": "workspace:^",

packages/wrangler/src/ai/fetcher.ts

@@ -0,0 +1,20 @@
+import { Response } from "miniflare";
+import { performApiFetch } from "../cfetch/internal";
+import { getAccountId } from "../user";
+import type { Request } from "miniflare";
+
+export async function AIFetcher(request: Request) {
+	const accountId = await getAccountId();
+
+	request.headers.delete("Host");
+	request.headers.delete("Content-Length");
+
+	const res = await performApiFetch(`/accounts/${accountId}/ai/run/proxy`, {
+		method: "POST",
+		headers: Object.fromEntries(request.headers.entries()),
+		body: request.body,
+		duplex: "half",
+	});
+
+	return new Response(res.body, { status: res.status });
+}

packages/wrangler/src/dev/miniflare.ts

@@ -2,6 +2,7 @@
 import { realpathSync } from "node:fs";
 import path from "node:path";
 import { Log, LogLevel, TypedEventTarget, Mutex, Miniflare } from "miniflare";
+import { AIFetcher } from "../ai/fetcher";
 import { ModuleTypeToRuleType } from "../deployment-bundle/module-collection";
 import { withSourceURLs } from "../deployment-bundle/source-url";
 import { getHttpsOptions } from "../https-options";
@@ -312,6 +313,10 @@
 				.join("\n"),
 	};
 
+	if (bindings.ai?.binding) {
+		config.serviceBindings[bindings.ai.binding] = AIFetcher;
+	}
+
 	const bindingOptions = {
 		bindings: bindings.vars,
 		textBlobBindings,
@@ -502,13 +507,7 @@
 		logger.warn("Miniflare 3 does not support CRON triggers yet, ignoring...");
 	}
 
-	if (config.bindings.ai) {
-		logger.warn(
-			"Workers AI is not currently supported in local mode. Please use --remote to work with it."
-		);
-	}
-
-	if (!config.bindings.ai && config.bindings.vectorize?.length) {
+	if (config.bindings.vectorize?.length) {
 		// TODO: add local support for Vectorize bindings (https://github.com/cloudflare/workers-sdk/issues/4360)
 		logger.warn(
 			"Vectorize bindings are not currently supported in local mode. Please use --remote if you are working with them."