assistant-cloud/src/tests/CloudRunReporter.test.ts

Cloud integration for assistant-ui

import { describe, expect, it, vi } from "vitest";
import type { AssistantCloud } from "../AssistantCloud";
import { CloudAPIError } from "../AssistantCloudAPI";
import { CloudRunReporter } from "../CloudRunReporter";
import { CloudResponseError } from "../cloudResponse";

const createCloud = (telemetry: AssistantCloud["telemetry"]) => {
  const report = vi.fn().mockResolvedValue({ run_id: "run_1" });
  const cloud = { telemetry, runs: { report } } as unknown as AssistantCloud;
  return { cloud, report };
};

describe("CloudRunReporter", () => {
  it("sends nothing while telemetry is disabled", async () => {
    const { cloud, report } = createCloud({ enabled: false });
    await new CloudRunReporter(cloud).report({
      threadId: "thread_1",
      status: "completed",
    });
    expect(report).not.toHaveBeenCalled();
  });

  it("stamps the cloud's environment, release and tags and applies beforeReport last", async () => {
    const beforeReport = vi.fn((report) => ({
      ...report,
      model_id: "gpt",
      cost_usd: 0.012,
      attributes: { tenant: "acme" },
      root_span_id: "0011223344556677",
    }));
    const { cloud, report } = createCloud({
      enabled: true,
      environment: "production",
      release: "1.2.3",
      tags: ["web"],
      beforeReport,
    });
    await new CloudRunReporter(cloud).report({
      threadId: "thread_1",
      status: "completed",
      outputText: "hi",
    });
    expect(beforeReport).toHaveBeenCalledWith(
      expect.objectContaining({
        thread_id: "thread_1",
        environment: "production",
        release: "1.2.3",
        tags: ["web"],
      }),
    );
    expect(report).toHaveBeenCalledWith(
      expect.objectContaining({
        thread_id: "thread_1",
        model_id: "gpt",
        cost_usd: 0.012,
        attributes: { tenant: "acme" },
        root_span_id: "0011223344556677",
      }),
    );
  });

  it("skips a report the hook vetoes without consuming its key", async () => {
    const { cloud, report } = createCloud({
      enabled: true,
      beforeReport: vi
        .fn()
        .mockReturnValueOnce(null)
        .mockImplementation((r) => r),
    });
    const reporter = new CloudRunReporter(cloud);
    await reporter.report({ threadId: "t", status: "completed" }, "t:m");
    await reporter.report({ threadId: "t", status: "completed" }, "t:m");
    expect(report).toHaveBeenCalledOnce();
  });

  it("reports a keyed run once and an unkeyed run every time", async () => {
    const { cloud, report } = createCloud({ enabled: true });
    const reporter = new CloudRunReporter(cloud);
    await reporter.report({ threadId: "t", status: "completed" }, "t:m");
    await reporter.report({ threadId: "t", status: "completed" }, "t:m");
    await reporter.report({ threadId: "t", status: "completed" });
    await reporter.report({ threadId: "t", status: "completed" });
    expect(report).toHaveBeenCalledTimes(3);
  });

  it("allows a keyed run to retry after rate limiting", async () => {
    const { cloud, report } = createCloud({ enabled: true });
    report.mockRejectedValueOnce(new CloudAPIError("rate limited", 429));
    const reporter = new CloudRunReporter(cloud);

    await reporter.report({ threadId: "t", status: "completed" }, "t:m");
    await reporter.report({ threadId: "t", status: "completed" }, "t:m");
    await reporter.report({ threadId: "t", status: "completed" }, "t:m");

    expect(report).toHaveBeenCalledTimes(2);
  });

  it.each([
    ["a transport failure", new Error("offline")],
    ["a request timeout", new CloudAPIError("timeout", 408)],
    ["a server error", new CloudAPIError("unavailable", 503)],
    [
      "a successful response with an invalid body",
      new CloudResponseError("invalid response"),
    ],
    ["a successful response with invalid JSON", new SyntaxError("invalid")],
    ["a permanent client error", new CloudAPIError("invalid report", 400)],
  ])("does not retry after %s", async (_name, error) => {
    const { cloud, report } = createCloud({ enabled: true });
    report.mockRejectedValueOnce(error);
    const reporter = new CloudRunReporter(cloud);

    await reporter.report({ threadId: "t", status: "completed" }, "t:m");
    await reporter.report({ threadId: "t", status: "completed" }, "t:m");

    expect(report).toHaveBeenCalledOnce();
  });

  it("swallows a failed send and reads the cloud through a getter", async () => {
    const { cloud } = createCloud({ enabled: true });
    (cloud.runs.report as ReturnType<typeof vi.fn>).mockRejectedValue(
      new Error("offline"),
    );
    await expect(
      new CloudRunReporter(() => cloud).report({
        threadId: "t",
        status: "error",
      }),
    ).resolves.toBeUndefined();
  });

  it("resolves when the hook throws and reports nothing", async () => {
    const { cloud, report } = createCloud({
      enabled: true,
      beforeReport: () => {
        throw new Error("hook");
      },
    });
    await expect(
      new CloudRunReporter(cloud).report({
        threadId: "t",
        status: "completed",
      }),
    ).resolves.toBeUndefined();
    expect(report).not.toHaveBeenCalled();
  });
});