Retry badge moderation when reasoning eats the whole budget

Source/Process.ts

@@ -46,6 +46,14 @@ function sleep(time: number) {
 }
 
 export const BadgeModerationModel = "@cf/zai-org/glm-4.7-flash";
+// Reasoning tokens come out of the same budget as the answer, so a badge the model
+// finds hard can spend the whole allowance deliberating and return no content at
+// all. Ordinary badges finish in a few hundred tokens; the headroom is only ever
+// billed on the inputs that need it.
+export const BadgeModerationMaxTokens = 2048;
+// One retry, because truncation is not a verdict and the model does not stop in the
+// same place twice.
+export const BadgeModerationAttempts = 2;
 const BadgeMaxGraphemes = 20;
 const BadgeEditsPerHour = 10;
 const BadgeQuotaWindow = 60 * 60 * 1000;
@@ -157,6 +165,13 @@ export function ReadModerationVerdict(Reply: any): { allowed: boolean, rule: num
   return {allowed: Payload.allowed, rule: Payload.rule};
 }
 
+// A reply that ran out of budget mid-thought says nothing about the badge, so it is
+// worth asking again. Anything else that fails to validate is the model answering
+// badly, and asking again would only spend another inference call to hear it twice.
+export function ModerationReplyTruncated(Reply: any): boolean {
+  return Reply?.choices?.[0]?.finish_reason === "length";
+}
+
 // 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.
@@ -1554,22 +1569,29 @@ export class Process {
         }
 
         let Verdict: { allowed: boolean, rule: number } | null = null;
-        try {
-          Verdict = ReadModerationVerdict(await this.AI.run(BadgeModerationModel, {
-            messages: [
-              {role: "system", content: BadgeModerationPrompt},
-              {role: "user", content: "<badge>" + Data["Content"] + "</badge>"}
-            ],
-            temperature: 0,
-            // The model reasons before answering; a small budget is spent entirely
-            // on reasoning and comes back with no content at all.
-            max_completion_tokens: 1024,
-            response_format: {type: "json_schema", json_schema: BadgeModerationSchema}
-          }));
-        } catch (Error) {
-          Output.Error("Badge moderation failed: " + Error + "\n" +
+        for (let Attempt = 0; Attempt < BadgeModerationAttempts; Attempt++) {
+          let Reply: any;
+          try {
+            Reply = await this.AI.run(BadgeModerationModel, {
+              messages: [
+                {role: "system", content: BadgeModerationPrompt},
+                {role: "user", content: "<badge>" + Data["Content"] + "</badge>"}
+              ],
+              temperature: 0,
+              max_completion_tokens: BadgeModerationMaxTokens,
+              response_format: {type: "json_schema", json_schema: BadgeModerationSchema}
+            });
+          } catch (Error) {
+            Output.Error("Badge moderation failed: " + Error + "\n" +
+              "Username: " + this.Username);
+            return new Result(false, "内容审核服务暂时不可用，请稍后重试");
+          }
+          Verdict = ReadModerationVerdict(Reply);
+          if (Verdict !== null || !ModerationReplyTruncated(Reply)) {
+            break;
+          }
+          Output.Warn("Badge moderation ran out of tokens before answering, retrying\n" +
             "Username: " + this.Username);
-          return new Result(false, "内容审核服务暂时不可用，请稍后重试");
         }
         if (Verdict === null) {
           Output.Error("Badge moderation returned an unusable verdict\n" +

test/process.test.js

@@ -564,8 +564,57 @@ test('EditBadge sends the badge to the moderation model, delimited', async () =>
     assert.strictEqual(seenModel, '@cf/zai-org/glm-4.7-flash');
     assert.strictEqual(seenBody.messages[1].content, '<badge>hello</badge>');
     assert.strictEqual(seenBody.temperature, 0);
-    // A small budget is spent entirely on reasoning and returns no content at all.
-    assert.ok(seenBody.max_completion_tokens >= 1024);
+    // Reasoning shares the completion budget, and a badge the model finds hard can
+    // spend more than 1024 tokens thinking before it writes any JSON.
+    assert.ok(seenBody.max_completion_tokens >= 2048);
+});
+
+// The shape the live model returns when reasoning eats the whole budget: no content,
+// and finish_reason "length" to say why.
+function truncated() {
+    return { choices: [{ finish_reason: 'length', message: { content: null } }] };
+}
+
+test('EditBadge retries when the model runs out of tokens before answering', async () => {
+    const replies = [truncated(), verdict(false, 1)];
+    const proc = createBadgeProcess({ ai: async () => replies.shift() });
+    const result = await proc.ProcessFunctions['EditBadge'](editArgs('陈开尔万岁'));
+    assert.strictEqual(proc.AI.run.mock.callCount(), 2);
+    assert.strictEqual(result.Success, false);
+    assert.strictEqual(result.Message, '标签内容包含不雅或粗俗用语，请修改后重试');
+});
+
+test('EditBadge keeps an allow verdict that arrives on the retry', async () => {
+    const replies = [truncated(), allow()];
+    const proc = createBadgeProcess({ ai: async () => replies.shift() });
+    const result = await proc.ProcessFunctions['EditBadge'](editArgs('hello'));
+    assert.ok(result.Success, result.Message);
+    assert.strictEqual(proc.AI.run.mock.callCount(), 2);
+});
+
+test('EditBadge fails closed when the retry is truncated too', async () => {
+    const proc = createBadgeProcess({ ai: async () => truncated() });
+    const result = await proc.ProcessFunctions['EditBadge'](editArgs('hello'));
+    assert.strictEqual(result.Success, false);
+    assert.strictEqual(result.Message, '内容审核服务暂时不可用，请稍后重试');
+    // Bounded: a badge that always truncates must not loop on the model.
+    assert.strictEqual(proc.AI.run.mock.callCount(), 2);
+});
+
+test('EditBadge does not retry a reply that is merely malformed', async () => {
+    const proc = createBadgeProcess({
+        ai: async () => ({ choices: [{ finish_reason: 'stop', message: { content: 'looks fine to me' } }] })
+    });
+    const result = await proc.ProcessFunctions['EditBadge'](editArgs('hello'));
+    assert.strictEqual(result.Success, false);
+    assert.strictEqual(result.Message, '内容审核服务暂时不可用，请稍后重试');
+    assert.strictEqual(proc.AI.run.mock.callCount(), 1);
+});
+
+test('EditBadge does not retry after the model throws', async () => {
+    const proc = createBadgeProcess({ ai: async () => { throw new Error('AI down'); } });
+    await proc.ProcessFunctions['EditBadge'](editArgs('hello'));
+    assert.strictEqual(proc.AI.run.mock.callCount(), 1);
 });
 
 test('EditBadge tells the user which rule the badge broke', async () => {

tools/check-badge.js

@@ -29,10 +29,13 @@ const os = require("os");
 const path = require("path");
 
 const {
+    BadgeModerationAttempts,
+    BadgeModerationMaxTokens,
     BadgeModerationModel,
     BadgeModerationPrompt,
     BadgeModerationSchema,
     BadgeRuleReasons,
+    ModerationReplyTruncated,
     ReadModerationVerdict,
 } = require("../Source/Process.ts");
 
@@ -66,7 +69,7 @@ function token() {
     return match ? match[1] : null;
 }
 
-async function moderate(content, bearer) {
+async function ask(content, bearer) {
     const response = await fetch(
         "https://api.cloudflare.com/client/v4/accounts/" + ACCOUNT_ID + "/ai/run/" + BadgeModerationModel,
         {
@@ -78,19 +81,32 @@ async function moderate(content, bearer) {
                     { role: "user", content: "<badge>" + content + "</badge>" },
                 ],
                 temperature: 0,
-                max_completion_tokens: 1024,
+                max_completion_tokens: BadgeModerationMaxTokens,
                 response_format: { type: "json_schema", json_schema: BadgeModerationSchema },
             }),
         }
     );
-    const body = await response.json();
-    if (!body.success) {
-        return { error: JSON.stringify(body.errors || body) };
+    return response.json();
+}
+
+// Same retry EditBadge does, so the neuron count printed below is what a real edit
+// of this text would cost, retries included.
+async function moderate(content, bearer) {
+    let neurons = 0;
+    let retries = 0;
+    for (let attempt = 0; attempt < BadgeModerationAttempts; attempt++) {
+        const body = await ask(content, bearer);
+        if (!body.success) {
+            return { error: JSON.stringify(body.errors || body), neurons };
+        }
+        neurons += body.result?.usage?.neurons || 0;
+        const verdict = ReadModerationVerdict(body.result);
+        if (verdict !== null || !ModerationReplyTruncated(body.result)) {
+            return { verdict, neurons, retries };
+        }
+        retries++;
     }
-    return {
-        verdict: ReadModerationVerdict(body.result),
-        neurons: body.result?.usage?.neurons,
-    };
+    return { verdict: null, neurons, retries };
 }
 
 async function main() {
@@ -123,19 +139,21 @@ async function main() {
             console.log("BLOCKED  " + JSON.stringify(content) + "  " + blocked + "  (no model call)");
             continue;
         }
-        const { verdict, neurons, error } = await moderate(content, bearer);
+        const { verdict, neurons, retries, error } = await moderate(content, bearer);
+        spent += neurons || 0;
         if (error) {
             console.log("ERROR    " + JSON.stringify(content) + "  " + error);
             continue;
         }
-        spent += neurons || 0;
+        const retried = retries > 0 ? "  (retried " + retries + "x after truncation)" : "";
         if (verdict === null) {
-            console.log("UNUSABLE " + JSON.stringify(content) + "  model reply failed validation, edit would fail closed");
+            console.log("UNUSABLE " + JSON.stringify(content) +
+                "  model reply failed validation, edit would fail closed" + retried);
         } else if (verdict.allowed) {
-            console.log("ALLOW    " + JSON.stringify(content));
+            console.log("ALLOW    " + JSON.stringify(content) + retried);
         } else {
             console.log("REJECT   " + JSON.stringify(content) +
-                "  rule " + verdict.rule + " — 标签内容" + BadgeRuleReasons[verdict.rule] + "，请修改后重试");
+                "  rule " + verdict.rule + " — 标签内容" + BadgeRuleReasons[verdict.rule] + "，请修改后重试" + retried);
         }
     }
     if (spent > 0) {