UHD_Media_Tool/apps/server/test/agent/spend-no-budget.test.ts

import { DEFAULT_SETTINGS, DEPRECATED_SETTINGS_FIELDS, type GlobalSettings } from "@umt/contracts";
import type { PiUpstreamRow } from "@umt/db";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { createAgentModule } from "../../src/modules/agent/index.ts";
import { agentSettings } from "../../src/modules/agent/runner.ts";
import { ORGANIZE_IDENTIFY_ACTOR } from "../../src/modules/agent/usage.ts";
import { createHarness, type Harness, pgAvailable } from "../auth-settings/harness.ts";
import { DEV_KEY, json } from "../helpers.ts";

// 智能体不设硬性预算上限(2026-10-06):每天次数 / 每天花费 / 单次花费 / 单次 token 四个上限不再拦截(旧值兼容读取、不生效),
// 只保留单次工具调用次数与单次最长时长两个防跑飞限制(缺省调宽);花费改为可见:/agent/usage 按会话类型统计今天与最近 7 天。

const upstream = { row: { provider: "fake", default_model: null, thinking_level: "off" } as unknown as PiUpstreamRow };

describe("预算上限移除,防跑飞限制保留", () => {
  it("旧设置里的 token / 费用上限不进运行参数;工具调用次数与时长照用", () => {
    const legacy = {
      ...DEFAULT_SETTINGS,
      agent: { ...DEFAULT_SETTINGS.agent, maxTokensPerRun: 1000, maxCostUsdPerRun: 0.01 },
      budget: { ...DEFAULT_SETTINGS.budget, agentMaxRunsPerDay: 0, agentMaxCostUsdPerDay: 0 },
    } as GlobalSettings;
    // biome-ignore lint/suspicious/noExplicitAny: 只取用到的上游字段
    const s = agentSettings(legacy, {}, true, upstream as any);
    expect(s).not.toHaveProperty("costBudget");
    expect(s.maxToolCalls).toBe(DEFAULT_SETTINGS.agent.maxToolCalls);
    expect(s.maxDurationMs).toBe(DEFAULT_SETTINGS.agent.maxDurationSec * 1000);
    // 调用方显式划拨(子会话、巡检预算)仍然生效
    // biome-ignore lint/suspicious/noExplicitAny: 同上
    const sub = agentSettings(legacy, { maxToolCalls: 3, maxCostUsd: 0.5 }, true, upstream as any);
    expect(sub.maxToolCalls).toBe(3);
    expect(sub.costBudget).toEqual({ maxCostUsd: 0.5 });
  });

  it("缺省调宽;四个上限登记为弃用字段", () => {
    expect(DEFAULT_SETTINGS.agent.maxToolCalls).toBeGreaterThanOrEqual(100);
    expect(DEFAULT_SETTINGS.agent.maxDurationSec).toBeGreaterThanOrEqual(1800);
    expect(DEFAULT_SETTINGS.agent.investigateMaxToolCalls).toBeGreaterThanOrEqual(20);
    expect(DEFAULT_SETTINGS.agent.autoIdentifyOrganize).toBe(true);
    expect(DEFAULT_SETTINGS.budget).not.toHaveProperty("agentMaxRunsPerDay");
    expect(DEPRECATED_SETTINGS_FIELDS).toEqual(
      expect.arrayContaining([
        "budget.agentMaxRunsPerDay",
        "budget.agentMaxCostUsdPerDay",
        "agent.maxTokensPerRun",
        "agent.maxCostUsdPerRun",
      ]),
    );
  });
});

const hasPg = await pgAvailable();

describe.skipIf(!hasPg)("花费统计 /agent/usage(真实 PG)", () => {
  let h: Harness;
  beforeAll(async () => {
    h = await createHarness({ extraModules: [createAgentModule()] });
  }, 120_000);
  afterAll(async () => {
    await h?.close();
  });

  it("今天与最近 7 天按会话类型统计次数、token 与费用;超过 7 天的不算", async () => {
    const day = 86_400_000;
    const now = Date.now();
    const run = async (
      trigger: string,
      at: number,
      usage: { totalTokens: number; costUsd: number },
      extra: Record<string, unknown> = {},
    ) =>
      (
        await h.db
          .insertInto("agent_run")
          .values({
            trigger,
            goal: "g",
            status: "succeeded",
            mode: "execute",
            usage: JSON.stringify(usage),
            started_at: new Date(at),
            created_at: new Date(at),
            ...extra,
          } as never)
          .returning("id")
          .executeTakeFirstOrThrow()
      ).id;
    const chat = await run("user", now, { totalTokens: 1000, costUsd: 0.1 });
    await run("subagent", now, { totalTokens: 500, costUsd: 0.05 }, { parent_run_id: chat });
    await run("event", now, { totalTokens: 200, costUsd: 0.02 }, { created_by: ORGANIZE_IDENTIFY_ACTOR });
    await run("event", now - 3 * day, { totalTokens: 300, costUsd: 0.03 });
    await run("schedule", now - 3 * day, { totalTokens: 400, costUsd: 0.04 });
    await run("user", now - 10 * day, { totalTokens: 9999, costUsd: 9 });

    const res = await h.app.request("/api/v1/agent/usage", { headers: { authorization: `Bearer ${DEV_KEY}` } });
    expect(res.status).toBe(200);
    const { today, last7d } = (await json(res)).data;
    expect(today.runs).toBe(3);
    expect(today.tokens).toBe(1700);
    expect(today.costUsd).toBeCloseTo(0.17);
    expect(today.byKind.map((k: { kind: string }) => k.kind)).toEqual(["chat", "organize_identify", "subagent"]);
    expect(last7d.runs).toBe(5);
    expect(last7d.costUsd).toBeCloseTo(0.24);
    expect(last7d.byKind).toEqual(
      expect.arrayContaining([
        { kind: "patrol", runs: 1, tokens: 400, costUsd: 0.04 },
        { kind: "investigate", runs: 1, tokens: 300, costUsd: 0.03 },
      ]),
    );
  });
});