Jyotisha/frontend/tests/consultation-workflow-retry.test.ts

import assert from "node:assert/strict";
import { readFileSync } from "node:fs";
import test from "node:test";
import { ConsultationWorkflowError } from "../src/mastra/consultation-workflow.ts";
import { CONSULTATION_QUEUE_LABEL } from "../src/lib/consultation-activity-labels.ts";
import {
  CONSULTATION_QUEUE_FAILED_MESSAGE,
  consultationRunFailureMessage,
  isConsultationRateLimitFailure,
  retryOnConsultationRateLimit,
} from "../src/lib/consultation-workflow-retry.ts";
import {
  emptyConsultationTimeline,
  reduceConsultationTimelineEvents,
} from "../src/lib/consultation-run-timeline.ts";

test("a 429 is retried once and the queue line is shown before success", async () => {
  let calls = 0;
  const queued: string[] = [];
  const sleeps: number[] = [];
  const result = await retryOnConsultationRateLimit({
    run: async () => {
      calls += 1;
      if (calls === 1) throw new ConsultationWorkflowError("workflow_rate_limited", "busy", 2);
      return "ok";
    },
    onQueue: () => { queued.push(CONSULTATION_QUEUE_LABEL); },
    sleep: async (ms) => { sleeps.push(ms); },
  });
  assert.equal(result, "ok");
  assert.equal(calls, 2);
  assert.deepEqual(sleeps, [2000]);
  assert.deepEqual(queued, ["人有点多，排一下队"]);
});

test("three 429s stop within two retries and 15 seconds", async () => {
  let calls = 0;
  let slept = 0;
  await assert.rejects(
    () => retryOnConsultationRateLimit({
      run: async () => {
        calls += 1;
        throw new ConsultationWorkflowError("workflow_rate_limited", "busy", 2);
      },
      sleep: async (ms) => { slept += ms; },
    }),
    (error: unknown) => {
      assert.equal(isConsultationRateLimitFailure(error), true);
      return true;
    },
  );
  assert.equal(calls, 3);
  assert.equal(slept, 4_000);
  assert.ok(slept <= 15_000);
  assert.equal(
    consultationRunFailureMessage(new ConsultationWorkflowError("workflow_rate_limited", "busy", 2), "其他失败"),
    CONSULTATION_QUEUE_FAILED_MESSAGE,
  );
  assert.equal(consultationRunFailureMessage(new Error("other"), "其他失败"), "其他失败");
});

test("Retry-After cannot push the total wait past 15 seconds", async () => {
  let calls = 0;
  let slept = 0;
  await assert.rejects(() => retryOnConsultationRateLimit({
    run: async () => {
      calls += 1;
      throw new ConsultationWorkflowError("workflow_rate_limited", "busy", 20);
    },
    sleep: async (ms) => { slept += ms; },
  }));
  assert.equal(calls, 2);
  assert.equal(slept, 15_000);
});

test("the queue sentence replaces the calculate row and the chart label comes back", () => {
  const calculating = reduceConsultationTimelineEvents([
    { type: "tool.started", callId: "tool-1", tool: "run-jyotish-consultation", label: "正在计算个人星盘" },
    { type: "activity", phase: "chart-calculation", label: "正在计算本命盘…" },
  ], emptyConsultationTimeline());
  const queued = reduceConsultationTimelineEvents([
    { type: "activity", phase: "compute-queue", label: CONSULTATION_QUEUE_LABEL },
  ], calculating);
  assert.equal(queued.rows.find((row) => row.kind === "calculate")?.label, "人有点多，排一下队");
  assert.equal(queued.rows.find((row) => row.kind === "calculate")?.status, "live");
  const resumed = reduceConsultationTimelineEvents([
    { type: "activity", phase: "chart-calculation", label: "正在计算本命盘…" },
  ], queued);
  assert.equal(resumed.rows.find((row) => row.kind === "calculate")?.label, "正在计算本命盘…");
});

test("a failed consultation still settles through the existing cancel path before the queue sentence", () => {
  const stream = readFileSync(new URL("../src/lib/stream-agent-response.ts", import.meta.url), "utf8");
  const consult = readFileSync(new URL("../src/app/api/consult/route.ts", import.meta.url), "utf8");
  const onErrorAt = stream.indexOf("await options.onError?.");
  const messageAt = stream.indexOf("consultationRunFailureMessage(error, runFailedMessage(code))");
  assert.ok(onErrorAt > 0 && messageAt > onErrorAt);
  assert.match(consult, /onError: \(error\) => settleRun\(\s*cancel/);
  assert.match(consult, /else await cancel\(\)/);
  assert.match(consult, /retryOnConsultationRateLimit/);
  assert.match(consult, /CONSULTATION_QUEUE_FAILED_MESSAGE/);
});