Skip to content

Commit 415520e

Browse files
penalosapenalosa
authored
Implement Mixed Mode proxy server & client (#9149)
* Implement proxy server * Use test tokens * Create red-trains-drive.md * Add auth option * lint * add turbo.json * stream logs * simplify auth * Implement AI Mixed Mode in Miniflare * pass through tokens * Update red-trains-drive.md * Update test-and-check.yml * Update turbo.json * 400
1 parent 02f0699 commit 415520e

File tree

11 files changed

+216
-26
lines changed

11 files changed

+216
-26
lines changed

‎.changeset/red-trains-drive.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"wrangler": patch
3+
---
4+
5+
Implement mixed mode proxy server & client

‎.github/workflows/test-and-check.yml‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -169,8 +169,8 @@ jobs:
169169
if: steps.changes.outputs.everything_but_markdown == 'true'
170170
run: pnpm run test:ci --concurrency 1 ${{ matrix.filter }}
171171
env:
172-
TMP_CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }}
173-
TMP_CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }}
172+
TEST_CLOUDFLARE_API_TOKEN: ${{ secrets.TEST_CLOUDFLARE_API_TOKEN }}
173+
TEST_CLOUDFLARE_ACCOUNT_ID: ${{ secrets.TEST_CLOUDFLARE_ACCOUNT_ID }}
174174
NODE_OPTIONS: "--max_old_space_size=8192"
175175
WRANGLER_LOG_PATH: ${{ runner.temp }}/wrangler-debug-logs/
176176
TEST_REPORT_PATH: ${{ runner.temp }}/test-report/index.html
@@ -226,8 +226,8 @@ jobs:
226226
if: steps.changes.outputs.everything_but_markdown == 'true'
227227
run: pnpm run test:ci --concurrency 1 ${{ matrix.filter }}
228228
env:
229-
TMP_CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }}
230-
TMP_CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }}
229+
TEST_CLOUDFLARE_API_TOKEN: ${{ secrets.TEST_CLOUDFLARE_API_TOKEN }}
230+
TEST_CLOUDFLARE_ACCOUNT_ID: ${{ secrets.TEST_CLOUDFLARE_ACCOUNT_ID }}
231231
NODE_OPTIONS: "--max_old_space_size=8192"
232232
WRANGLER_LOG_PATH: ${{ runner.temp }}/wrangler-debug-logs/
233233
TEST_REPORT_PATH: ${{ runner.temp }}/test-report/index.html

‎fixtures/mixed-mode-node-test/package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
},
1111
"devDependencies": {
1212
"@cloudflare/workers-tsconfig": "workspace:*",
13+
"miniflare": "workspace:*",
1314
"wrangler": "workspace:*"
1415
},
1516
"volta": {

‎fixtures/mixed-mode-node-test/tests/startMixedModeSession.test.js‎

Lines changed: 64 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,76 @@
11
import assert from "node:assert";
22
import test, { describe } from "node:test";
3+
import { Miniflare } from "miniflare";
34
import { experimental_startMixedModeSession } from "wrangler";
45

5-
describe("startMixedModeSession", () => {
6-
test("no-op mixed-mode proxyServerWorker", async (t) => {
7-
if (
8-
!process.env.CLOUDFLARE_ACCOUNT_ID ||
9-
!process.env.CLOUDFLARE_API_TOKEN
10-
) {
11-
return t.skip();
12-
}
6+
process.env.CLOUDFLARE_ACCOUNT_ID = process.env.TEST_CLOUDFLARE_ACCOUNT_ID;
7+
process.env.CLOUDFLARE_API_TOKEN = process.env.TEST_CLOUDFLARE_API_TOKEN;
138

14-
const mixedModeSession = await experimental_startMixedModeSession({});
9+
describe("startMixedModeSession", () => {
10+
test("simple AI request to the proxyServerWorker", async () => {
11+
const mixedModeSession = await experimental_startMixedModeSession({
12+
AI: {
13+
type: "ai",
14+
},
15+
});
1516
const proxyServerUrl =
1617
mixedModeSession.mixedModeConnectionString.toString();
1718
assert.match(proxyServerUrl, /http:\/\/localhost:\d{4,5}\//);
18-
assert.strictEqual(
19-
await (await fetch(proxyServerUrl)).text(),
20-
"no-op mixed-mode proxyServerWorker"
19+
assert.match(
20+
await (
21+
await fetch(proxyServerUrl, {
22+
headers: {
23+
"MF-Binding": "AI",
24+
"MF-URL": "https://workers-binding.ai/ai-api/models/search",
25+
},
26+
})
27+
).text(),
28+
// Assert the catalog _at least_ contains a LLama model
29+
/Llama/
30+
);
31+
await mixedModeSession.ready;
32+
await mixedModeSession.dispose();
33+
});
34+
test("AI mixed mode binding", async () => {
35+
const mixedModeSession = await experimental_startMixedModeSession({
36+
AI: {
37+
type: "ai",
38+
},
39+
});
40+
41+
const mf = new Miniflare({
42+
compatibilityDate: "2025-01-01",
43+
modules: true,
44+
script: /* javascript */ `
45+
export default {
46+
async fetch(request, env) {
47+
const messages = [
48+
{
49+
role: "user",
50+
// Doing snapshot testing against AI responses can be flaky, but this prompt generates the same output relatively reliably
51+
content: "Respond with the exact text 'This is a response from Workers AI.'. Do not include any other text",
52+
},
53+
];
54+
55+
const content = await env.AI.run("@hf/thebloke/zephyr-7b-beta-awq", {
56+
messages,
57+
});
58+
59+
return new Response(content.response);
60+
}
61+
}
62+
`,
63+
ai: {
64+
binding: "AI",
65+
mixedModeConnectionString: mixedModeSession.mixedModeConnectionString,
66+
},
67+
});
68+
assert.match(
69+
await (await mf.dispatchFetch("http://example.com")).text(),
70+
/This is a response from Workers AI/
2171
);
72+
await mf.dispose();
73+
2274
await mixedModeSession.ready;
2375
await mixedModeSession.dispose();
2476
});

‎fixtures/mixed-mode-node-test/turbo.json‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
{
2+
"$schema": "http://turbo.build/schema.json",
3+
"extends": ["//"],
4+
"tasks": {
5+
"test:ci": {
6+
"env": ["TEST_CLOUDFLARE_ACCOUNT_ID", "TEST_CLOUDFLARE_API_TOKEN"]
7+
}
8+
}
9+
}

‎packages/miniflare/src/plugins/ai/index.ts‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
import assert from "node:assert";
2+
import SCRIPT_MIXED_MODE_CLIENT from "worker:shared/mixed-mode-client";
3+
import { z } from "zod";
4+
import { MixedModeConnectionString, Plugin, ProxyNodeBinding } from "../shared";
5+
6+
const AISchema = z.object({
7+
binding: z.string(),
8+
mixedModeConnectionString: z.custom<MixedModeConnectionString>(),
9+
});
10+
11+
export const AIOptionsSchema = z.object({
12+
ai: AISchema.optional(),
13+
});
14+
15+
export const AI_PLUGIN_NAME = "ai";
16+
17+
export const AI_PLUGIN: Plugin<typeof AIOptionsSchema> = {
18+
options: AIOptionsSchema,
19+
async getBindings(options) {
20+
if (!options.ai) {
21+
return [];
22+
}
23+
24+
assert(
25+
options.ai.mixedModeConnectionString,
26+
"Workers AI only supports Mixed Mode"
27+
);
28+
29+
return [
30+
{
31+
name: options.ai.binding,
32+
wrapped: {
33+
moduleName: "cloudflare-internal:ai-api",
34+
innerBindings: [
35+
{
36+
name: "fetcher",
37+
service: { name: `${AI_PLUGIN_NAME}:${options.ai.binding}` },
38+
},
39+
],
40+
},
41+
},
42+
];
43+
},
44+
getNodeBindings(options: z.infer<typeof AIOptionsSchema>) {
45+
if (!options.ai) {
46+
return {};
47+
}
48+
return {
49+
[options.ai.binding]: new ProxyNodeBinding(),
50+
};
51+
},
52+
async getServices({ options }) {
53+
if (!options.ai) {
54+
return [];
55+
}
56+
57+
return [
58+
{
59+
name: `${AI_PLUGIN_NAME}:${options.ai.binding}`,
60+
worker: {
61+
compatibilityDate: "2025-01-01",
62+
modules: [
63+
{
64+
name: "index.worker.js",
65+
esModule: SCRIPT_MIXED_MODE_CLIENT(),
66+
},
67+
],
68+
bindings: [
69+
{
70+
name: "mixedModeConnectionString",
71+
text: options.ai.mixedModeConnectionString.href,
72+
},
73+
{
74+
name: "binding",
75+
text: options.ai.binding,
76+
},
77+
],
78+
},
79+
},
80+
];
81+
},
82+
};

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { z } from "zod";
22
import { ValueOf } from "../workers";
3+
import { AI_PLUGIN, AI_PLUGIN_NAME } from "./ai";
34
import {
45
ANALYTICS_ENGINE_PLUGIN,
56
ANALYTICS_ENGINE_PLUGIN_NAME,
@@ -36,6 +37,7 @@ export const PLUGINS = {
3637
[SECRET_STORE_PLUGIN_NAME]: SECRET_STORE_PLUGIN,
3738
[EMAIL_PLUGIN_NAME]: EMAIL_PLUGIN,
3839
[ANALYTICS_ENGINE_PLUGIN_NAME]: ANALYTICS_ENGINE_PLUGIN,
40+
[AI_PLUGIN_NAME]: AI_PLUGIN,
3941
};
4042
export type Plugins = typeof PLUGINS;
4143

@@ -88,7 +90,8 @@ export type WorkerOptions = z.input<typeof CORE_PLUGIN.options> &
8890
z.input<typeof WORKFLOWS_PLUGIN.options> &
8991
z.input<typeof PIPELINE_PLUGIN.options> &
9092
z.input<typeof SECRET_STORE_PLUGIN.options> &
91-
z.input<typeof ANALYTICS_ENGINE_PLUGIN.options>;
93+
z.input<typeof ANALYTICS_ENGINE_PLUGIN.options> &
94+
z.input<typeof AI_PLUGIN.options>;
9295

9396
export type SharedOptions = z.input<typeof CORE_PLUGIN.sharedOptions> &
9497
z.input<typeof CACHE_PLUGIN.sharedOptions> &
@@ -151,3 +154,4 @@ export * from "./pipelines";
151154
export * from "./secret-store";
152155
export * from "./email";
153156
export * from "./analytics-engine";
157+
export * from "./ai";

‎packages/miniflare/src/workers/shared/mixed-mode-client.worker.ts‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
export default {
2+
async fetch(request, env) {
3+
const proxiedHeaders = new Headers();
4+
for (const [name, value] of request.headers) {
5+
proxiedHeaders.set(`MF-Header-${name}`, value);
6+
}
7+
proxiedHeaders.set("MF-URL", request.url);
8+
proxiedHeaders.set("MF-Binding", env.binding);
9+
const req = new Request(request, {
10+
headers: proxiedHeaders,
11+
});
12+
13+
return fetch(env.mixedModeConnectionString, req);
14+
},
15+
} satisfies ExportedHandler<{
16+
mixedModeConnectionString: string;
17+
binding: string;
18+
}>;

‎packages/wrangler/src/api/mixedMode/index.ts‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
import path from "node:path";
22
import { getBasePath } from "../../paths";
3-
import { requireApiToken, requireAuth } from "../../user";
43
import { startWorker } from "../startDevWorker";
54
import type { StartDevWorkerInput, Worker } from "../startDevWorker/types";
65
import type { MixedModeConnectionString } from "miniflare";
@@ -13,7 +12,10 @@ type MixedModeSession = Pick<Worker, "ready" | "dispose"> & {
1312
};
1413

1514
export async function startMixedModeSession(
16-
bindings: BindingsOpt
15+
bindings: BindingsOpt,
16+
options?: {
17+
auth: NonNullable<StartDevWorkerInput["dev"]>["auth"];
18+
}
1719
): Promise<MixedModeSession> {
1820
const proxyServerWorkerWranglerConfig = path.resolve(
1921
getBasePath(),
@@ -24,10 +26,7 @@ export async function startMixedModeSession(
2426
config: proxyServerWorkerWranglerConfig,
2527
dev: {
2628
remote: true,
27-
auth: {
28-
accountId: await requireAuth({}),
29-
apiToken: requireApiToken(),
30-
},
29+
auth: options?.auth,
3130
},
3231
bindings,
3332
});

‎packages/wrangler/templates/mixedMode/proxyServerWorker/index.ts‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,22 @@
11
export default {
2-
fetch() {
3-
return new Response("no-op mixed-mode proxyServerWorker");
2+
async fetch(request, env) {
3+
const targetBinding = request.headers.get("MF-Binding");
4+
5+
if (targetBinding) {
6+
const originalHeaders = new Headers();
7+
for (const [name, value] of request.headers) {
8+
if (name.startsWith("mf-header-")) {
9+
originalHeaders.set(name.slice("mf-header-".length), value);
10+
}
11+
}
12+
return env[targetBinding].fetch(
13+
request.headers.get("MF-URL")!,
14+
new Request(request, {
15+
redirect: "manual",
16+
headers: originalHeaders,
17+
})
18+
);
19+
}
20+
return new Response("Provide a binding", { status: 400 });
421
},
5-
};
22+
} satisfies ExportedHandler<Record<string, Fetcher>>;

0 commit comments

Comments
 (0)