Call Workers AI over HTTP instead of the ai binding

Source/Process.ts

@@ -35,7 +35,8 @@ interface Environment {
   CaptchaSecretKey: string;
   DB: D1Database;
   logdb: AnalyticsEngineDataset;
-  AI: any;
+  AI?: any;
+  AI_API_TOKEN?: string;
   NOTIFICATIONS: DurableObjectNamespace;
   NOTIFICATION_PUSH_TOKEN: string;
 }
@@ -172,6 +173,31 @@ export function ModerationReplyTruncated(Reply: any): boolean {
   return Reply?.choices?.[0]?.finish_reason === "length";
 }
 
+// Workers AI over the REST API, shaped like the binding's `run` so callers do not
+// care which one they got. Cloudflare has been failing to attach the `ai` binding
+// to new versions of this worker (error 10021), so it is no longer in wrangler.toml;
+// the REST API serves the same models and returns the binding's output under
+// `result`. https://github.com/cloudflare/workers-sdk/issues/15973
+export function WorkersAIOverHTTP(AccountID: string, Token: string) {
+  return {
+    run: async (Model: string, Inputs: object): Promise<any> => {
+      const Response = await fetch(`https://api.cloudflare.com/client/v4/accounts/${AccountID}/ai/run/${Model}`, {
+        method: "POST",
+        headers: {
+          "Authorization": `Bearer ${Token}`,
+          "Content-Type": "application/json"
+        },
+        body: JSON.stringify(Inputs)
+      });
+      const Body: any = await Response.json().catch(() => null);
+      if (!Response.ok || !Body?.success) {
+        throw new Error("Workers AI returned HTTP " + Response.status + ": " + JSON.stringify(Body?.errors ?? null));
+      }
+      return Body.result;
+    }
+  };
+}
+
 // The KV key holding the list of problems that have a std answer. It is a
 // cache of `SELECT problem_id FROM std_answer`, kept so that GetStdList - a
 // hot read - costs no database rows.
@@ -1855,7 +1881,12 @@ export class Process {
 
   constructor(RequestData: Request, Environment: Environment) {
     this.XMOJDatabase = new Database(Environment.DB);
-    this.AI = Environment.AI;
+    // Prefer the binding if it is ever restored; otherwise go over HTTP, with a
+    // dedicated AI_API_TOKEN if one is set and the analytics API_TOKEN (granted
+    // Workers AI Read) if not. With no AI at all, `run` throws inside the
+    // moderation try block and the edit is refused.
+    const AIToken = Environment.AI_API_TOKEN || Environment.API_TOKEN;
+    this.AI = Environment.AI ?? (AIToken ? WorkersAIOverHTTP(Environment.ACCOUNT_ID, AIToken) : null);
     this.kv = Environment.kv;
     this.logs = Environment.logdb;
     this.notifications = Environment.NOTIFICATIONS;

Source/index.ts

@@ -29,7 +29,8 @@ interface Environment {
   CaptchaSecretKey: string;
   DB: D1Database;
   logdb: AnalyticsEngineDataset;
-  AI: Ai;
+  AI?: Ai;
+  AI_API_TOKEN?: string;
   NOTIFICATIONS: DurableObjectNamespace;
   NOTIFICATION_PUSH_TOKEN: string;
 }

test/process.test.js

@@ -1,6 +1,6 @@
 const test = require('node:test');
 const assert = require('node:assert');
-const { Process, RebuildStdList } = require('../Source/Process.ts');
+const { Process, RebuildStdList, WorkersAIOverHTTP } = require('../Source/Process.ts');
 const { Result } = require('../Source/Result.ts');
 
 function createProcess(mocks = {}) {
@@ -1211,3 +1211,87 @@ test('UploadStd rebuilds when the cache key is unset entirely', async () => {
 
     assert.strictEqual(kv.store.std_list, '1234');
 });
+
+// ---------------------------------------------------------------------------
+// Workers AI over HTTP
+// ---------------------------------------------------------------------------
+
+function processWithEnv(extra) {
+    const env = Object.assign({ ACCOUNT_ID: 'test-account-id', DB: { withSession() { return this; } } }, extra);
+    return new Process(new Request('https://example.com'), env);
+}
+
+test('WorkersAIOverHTTP posts the inputs to the account run endpoint and returns result', async (t) => {
+    let seenUrl = null, seenInit = null;
+    t.mock.method(global, 'fetch', async (url, init) => {
+        seenUrl = url;
+        seenInit = init;
+        return Response.json({ success: true, result: verdict(true, 0), errors: [] });
+    });
+    const reply = await WorkersAIOverHTTP('acct', 'tok').run('@cf/zai-org/glm-4.7-flash', { temperature: 0 });
+    assert.strictEqual(seenUrl, 'https://api.cloudflare.com/client/v4/accounts/acct/ai/run/@cf/zai-org/glm-4.7-flash');
+    assert.strictEqual(seenInit.method, 'POST');
+    assert.strictEqual(seenInit.headers['Authorization'], 'Bearer tok');
+    assert.deepStrictEqual(JSON.parse(seenInit.body), { temperature: 0 });
+    assert.deepStrictEqual(reply, verdict(true, 0));
+});
+
+test('WorkersAIOverHTTP throws on an HTTP error', async (t) => {
+    t.mock.method(global, 'fetch', async () => Response.json({ success: false, errors: [{ code: 10000 }] }, { status: 401 }));
+    await assert.rejects(WorkersAIOverHTTP('acct', 'tok').run('m', {}), /HTTP 401/);
+});
+
+test('WorkersAIOverHTTP throws when the API reports failure with 200', async (t) => {
+    t.mock.method(global, 'fetch', async () => Response.json({ success: false, errors: [] }));
+    await assert.rejects(WorkersAIOverHTTP('acct', 'tok').run('m', {}), /HTTP 200/);
+});
+
+test('WorkersAIOverHTTP throws on a non-JSON body', async (t) => {
+    t.mock.method(global, 'fetch', async () => new Response('<html>', { status: 502 }));
+    await assert.rejects(WorkersAIOverHTTP('acct', 'tok').run('m', {}), /HTTP 502/);
+});
+
+test('Process prefers the AI binding when it exists', () => {
+    const binding = { run: async () => ({}) };
+    const proc = processWithEnv({ AI: binding, AI_API_TOKEN: 'tok' });
+    assert.strictEqual(proc.AI, binding);
+});
+
+test('Process falls back to HTTP when only the token is set', async (t) => {
+    let seenUrl = null;
+    t.mock.method(global, 'fetch', async (url) => {
+        seenUrl = url;
+        return Response.json({ success: true, result: {} });
+    });
+    const proc = processWithEnv({ AI_API_TOKEN: 'tok' });
+    await proc.AI.run('m', {});
+    assert.strictEqual(seenUrl, 'https://api.cloudflare.com/client/v4/accounts/test-account-id/ai/run/m');
+});
+
+test('Process prefers AI_API_TOKEN over API_TOKEN', async (t) => {
+    let seenAuth = null;
+    t.mock.method(global, 'fetch', async (url, init) => {
+        seenAuth = init.headers['Authorization'];
+        return Response.json({ success: true, result: {} });
+    });
+    await processWithEnv({ API_TOKEN: 'analytics', AI_API_TOKEN: 'ai' }).AI.run('m', {});
+    assert.strictEqual(seenAuth, 'Bearer ai');
+});
+
+test('Process falls back to API_TOKEN for Workers AI', async (t) => {
+    let seenAuth = null;
+    t.mock.method(global, 'fetch', async (url, init) => {
+        seenAuth = init.headers['Authorization'];
+        return Response.json({ success: true, result: {} });
+    });
+    await processWithEnv({ API_TOKEN: 'analytics' }).AI.run('m', {});
+    assert.strictEqual(seenAuth, 'Bearer analytics');
+});
+
+test('EditBadge refuses the edit when there is no AI at all', async () => {
+    const proc = createBadgeProcess();
+    proc.AI = processWithEnv({}).AI;
+    const result = await proc.ProcessFunctions['EditBadge'](editArgs('hello'));
+    assert.strictEqual(result.Success, false);
+    assert.strictEqual(result.Message, '内容审核服务暂时不可用，请稍后重试');
+});

wrangler.toml

@@ -17,9 +17,6 @@ crons = ["0 0 * * *"]
 [[analytics_engine_datasets]]
 binding = "logdb"
 
-[ai]
-binding = "AI"
-
 [observability]
 enabled = true