claude-code-hub/tests/unit/proxy/session.test.ts

import { describe, expect, it, vi } from "vitest";
import type { ModelPrice, ModelPriceData } from "@/types/model-price";
import type { SystemSettings } from "@/types/system-config";
import type { Provider } from "@/types/provider";

vi.mock("@/repository/model-price", () => ({
  findLatestPriceByModel: vi.fn(),
}));

vi.mock("@/repository/system-config", () => ({
  getSystemSettings: vi.fn(),
}));

import { ProxySession } from "@/app/v1/_lib/proxy/session";
import { findLatestPriceByModel } from "@/repository/model-price";
import { getSystemSettings } from "@/repository/system-config";

function makeSystemSettings(
  billingModelSource: SystemSettings["billingModelSource"]
): SystemSettings {
  const now = new Date();
  return {
    id: 1,
    siteTitle: "test",
    allowGlobalUsageView: false,
    currencyDisplay: "USD",
    billingModelSource,
    timezone: null,
    enableAutoCleanup: false,
    cleanupRetentionDays: 30,
    cleanupSchedule: "0 2 * * *",
    cleanupBatchSize: 10000,
    enableClientVersionCheck: false,
    verboseProviderError: false,
    enableHttp2: false,
    interceptAnthropicWarmupRequests: false,
    enableResponseFixer: true,
    responseFixerConfig: {
      fixTruncatedJson: true,
      fixSseFormat: true,
      fixEncoding: true,
      maxJsonDepth: 200,
      maxFixSize: 1024 * 1024,
    },
    createdAt: now,
    updatedAt: now,
  };
}

function makePriceRecord(modelName: string, priceData: ModelPriceData): ModelPrice {
  return {
    id: 1,
    modelName,
    priceData,
    createdAt: new Date(),
    updatedAt: new Date(),
  };
}

function createSession({
  originalModel,
  redirectedModel,
  requestUrl,
  requestMessage,
}: {
  originalModel?: string | null;
  redirectedModel?: string | null;
  requestUrl?: URL;
  requestMessage?: Record<string, unknown>;
}): ProxySession {
  const session = new (
    ProxySession as unknown as {
      new (init: {
        startTime: number;
        method: string;
        requestUrl: URL;
        headers: Headers;
        headerLog: string;
        request: { message: Record<string, unknown>; log: string; model: string | null };
        userAgent: string | null;
        context: unknown;
        clientAbortSignal: AbortSignal | null;
      }): ProxySession;
    }
  )({
    startTime: Date.now(),
    method: "POST",
    requestUrl: requestUrl ?? new URL("http://localhost/v1/messages"),
    headers: new Headers(),
    headerLog: "",
    request: { message: requestMessage ?? {}, log: "(test)", model: redirectedModel ?? null },
    userAgent: null,
    context: {},
    clientAbortSignal: null,
  });

  if (originalModel !== undefined) {
    session.setOriginalModel(originalModel);
  }

  return session;
}

describe("ProxySession.getCachedPriceDataByBillingSource", () => {
  it("配置 = original 时应优先使用原始模型", async () => {
    const originalPriceData: ModelPriceData = { input_cost_per_token: 1, output_cost_per_token: 2 };
    const redirectedPriceData: ModelPriceData = {
      input_cost_per_token: 3,
      output_cost_per_token: 4,
    };

    vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("original"));
    vi.mocked(findLatestPriceByModel).mockImplementation(async (modelName: string) => {
      if (modelName === "original-model") {
        return makePriceRecord(modelName, originalPriceData);
      }
      if (modelName === "redirected-model") {
        return makePriceRecord(modelName, redirectedPriceData);
      }
      return null;
    });

    const session = createSession({
      originalModel: "original-model",
      redirectedModel: "redirected-model",
    });

    const result = await session.getCachedPriceDataByBillingSource();
    expect(result).toEqual(originalPriceData);
    expect(findLatestPriceByModel).toHaveBeenCalledTimes(1);
    expect(findLatestPriceByModel).toHaveBeenCalledWith("original-model");
  });

  it("配置 = redirected 时应优先使用重定向后模型", async () => {
    const originalPriceData: ModelPriceData = { input_cost_per_token: 1, output_cost_per_token: 2 };
    const redirectedPriceData: ModelPriceData = {
      input_cost_per_token: 3,
      output_cost_per_token: 4,
    };

    vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected"));
    vi.mocked(findLatestPriceByModel).mockImplementation(async (modelName: string) => {
      if (modelName === "original-model") {
        return makePriceRecord(modelName, originalPriceData);
      }
      if (modelName === "redirected-model") {
        return makePriceRecord(modelName, redirectedPriceData);
      }
      return null;
    });

    const session = createSession({
      originalModel: "original-model",
      redirectedModel: "redirected-model",
    });

    const result = await session.getCachedPriceDataByBillingSource();
    expect(result).toEqual(redirectedPriceData);
    expect(findLatestPriceByModel).toHaveBeenCalledTimes(1);
    expect(findLatestPriceByModel).toHaveBeenCalledWith("redirected-model");
  });

  it("应忽略空 priceData 并回退到备选模型", async () => {
    const redirectedPriceData: ModelPriceData = {
      input_cost_per_token: 3,
      output_cost_per_token: 4,
    };

    vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("original"));
    vi.mocked(findLatestPriceByModel)
      .mockResolvedValueOnce(makePriceRecord("original-model", {}))
      .mockResolvedValueOnce(makePriceRecord("redirected-model", redirectedPriceData));

    const session = createSession({
      originalModel: "original-model",
      redirectedModel: "redirected-model",
    });

    const result = await session.getCachedPriceDataByBillingSource();
    expect(result).toEqual(redirectedPriceData);
    expect(findLatestPriceByModel).toHaveBeenCalledTimes(2);
    expect(findLatestPriceByModel).toHaveBeenNthCalledWith(1, "original-model");
    expect(findLatestPriceByModel).toHaveBeenNthCalledWith(2, "redirected-model");
  });

  it("应在主模型无价格时回退到备选模型", async () => {
    const redirectedPriceData: ModelPriceData = {
      input_cost_per_token: 3,
      output_cost_per_token: 4,
    };

    vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("original"));
    vi.mocked(findLatestPriceByModel)
      .mockResolvedValueOnce(null)
      .mockResolvedValueOnce(makePriceRecord("redirected-model", redirectedPriceData));

    const session = createSession({
      originalModel: "original-model",
      redirectedModel: "redirected-model",
    });

    const result = await session.getCachedPriceDataByBillingSource();
    expect(result).toEqual(redirectedPriceData);
    expect(findLatestPriceByModel).toHaveBeenCalledTimes(2);
    expect(findLatestPriceByModel).toHaveBeenNthCalledWith(1, "original-model");
    expect(findLatestPriceByModel).toHaveBeenNthCalledWith(2, "redirected-model");
  });

  it("应在 getSystemSettings 失败时回退到 redirected", async () => {
    const redirectedPriceData: ModelPriceData = {
      input_cost_per_token: 3,
      output_cost_per_token: 4,
    };

    vi.mocked(getSystemSettings).mockRejectedValue(new Error("DB error"));
    vi.mocked(findLatestPriceByModel).mockResolvedValue(
      makePriceRecord("redirected-model", redirectedPriceData)
    );

    const session = createSession({
      originalModel: "original-model",
      redirectedModel: "redirected-model",
    });

    const result = await session.getCachedPriceDataByBillingSource();
    expect(result).toEqual(redirectedPriceData);
    expect(getSystemSettings).toHaveBeenCalledTimes(1);
    expect(findLatestPriceByModel).toHaveBeenCalledTimes(1);
    expect(findLatestPriceByModel).toHaveBeenCalledWith("redirected-model");

    const internal = session as unknown as { cachedBillingModelSource?: unknown };
    expect(internal.cachedBillingModelSource).toBe("redirected");
  });

  it("应在 billingModelSource 非法时回退到 redirected", async () => {
    const redirectedPriceData: ModelPriceData = {
      input_cost_per_token: 3,
      output_cost_per_token: 4,
    };

    vi.mocked(getSystemSettings).mockResolvedValue({
      ...makeSystemSettings("redirected"),
      billingModelSource: "invalid" as any,
    } as any);
    vi.mocked(findLatestPriceByModel).mockResolvedValue(
      makePriceRecord("redirected-model", redirectedPriceData)
    );

    const session = createSession({
      originalModel: "original-model",
      redirectedModel: "redirected-model",
    });

    const result = await session.getCachedPriceDataByBillingSource();
    expect(result).toEqual(redirectedPriceData);
    expect(findLatestPriceByModel).toHaveBeenCalledTimes(1);
    expect(findLatestPriceByModel).toHaveBeenCalledWith("redirected-model");

    const internal = session as unknown as { cachedBillingModelSource?: unknown };
    expect(internal.cachedBillingModelSource).toBe("redirected");
  });

  it("当原始模型等于重定向模型时应避免重复查询", async () => {
    vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("original"));
    vi.mocked(findLatestPriceByModel).mockResolvedValue(null);

    const session = createSession({
      originalModel: "same-model",
      redirectedModel: "same-model",
    });

    const result = await session.getCachedPriceDataByBillingSource();
    expect(result).toBeNull();
    expect(findLatestPriceByModel).toHaveBeenCalledTimes(1);
    expect(findLatestPriceByModel).toHaveBeenCalledWith("same-model");
  });

  it("并发调用时应只读取一次配置", async () => {
    const priceData: ModelPriceData = { input_cost_per_token: 1, output_cost_per_token: 2 };

    vi.mocked(getSystemSettings).mockImplementation(async () => {
      await new Promise((resolve) => setTimeout(resolve, 10));
      return makeSystemSettings("redirected");
    });
    vi.mocked(findLatestPriceByModel).mockResolvedValue(
      makePriceRecord("redirected-model", priceData)
    );

    const session = createSession({
      originalModel: "original-model",
      redirectedModel: "redirected-model",
    });

    const p1 = session.getCachedPriceDataByBillingSource();
    const p2 = session.getCachedPriceDataByBillingSource();
    await Promise.all([p1, p2]);

    expect(getSystemSettings).toHaveBeenCalledTimes(1);
  });

  it("应缓存配置避免重复读取", async () => {
    const priceData: ModelPriceData = { input_cost_per_token: 1, output_cost_per_token: 2 };

    vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected"));
    vi.mocked(findLatestPriceByModel).mockResolvedValue(
      makePriceRecord("redirected-model", priceData)
    );

    const session = createSession({
      originalModel: "original-model",
      redirectedModel: "redirected-model",
    });

    await session.getCachedPriceDataByBillingSource();
    await session.getCachedPriceDataByBillingSource();

    expect(getSystemSettings).toHaveBeenCalledTimes(1);
  });

  it("应缓存价格数据避免重复查询", async () => {
    const priceData: ModelPriceData = { input_cost_per_token: 1, output_cost_per_token: 2 };

    vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected"));
    vi.mocked(findLatestPriceByModel).mockResolvedValue(
      makePriceRecord("redirected-model", priceData)
    );

    const session = createSession({
      originalModel: "original-model",
      redirectedModel: "redirected-model",
    });

    await session.getCachedPriceDataByBillingSource();
    await session.getCachedPriceDataByBillingSource();

    expect(findLatestPriceByModel).toHaveBeenCalledTimes(1);
  });

  it("应在无模型时返回 null", async () => {
    vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected"));

    const session = createSession({ redirectedModel: null });
    const result = await session.getCachedPriceDataByBillingSource();

    expect(result).toBeNull();
    expect(getSystemSettings).not.toHaveBeenCalled();
    expect(findLatestPriceByModel).not.toHaveBeenCalled();
  });
});

function createSessionForHeaders(headers: Headers): ProxySession {
  // 使用 ProxySession 的内部构造方法创建测试实例
  const testSession = ProxySession.fromContext as any;
  const session = Object.create(ProxySession.prototype);

  Object.assign(session, {
    startTime: Date.now(),
    method: "POST",
    requestUrl: new URL("https://example.com/v1/messages"),
    headers,
    originalHeaders: new Headers(headers), // 同步更新 originalHeaders
    headerLog: JSON.stringify(Object.fromEntries(headers.entries())),
    request: { message: {}, log: "" },
    userAgent: headers.get("user-agent"),
    context: null,
    clientAbortSignal: null,
    userName: "test-user",
    authState: null,
    provider: null,
    messageContext: null,
    sessionId: null,
    requestSequence: 1,
    originalFormat: "claude",
    providerType: null,
    originalModelName: null,
    originalUrlPathname: null,
    providerChain: [],
    cacheTtlResolved: null,
    context1mApplied: false,
    cachedPriceData: undefined,
    cachedBillingModelSource: undefined,
  });

  return session;
}

describe("ProxySession - isHeaderModified", () => {
  it("应该检测到被修改的 header", () => {
    const headers = new Headers([["user-agent", "original"]]);
    const session = createSessionForHeaders(headers);

    session.headers.set("user-agent", "modified");

    expect(session.isHeaderModified("user-agent")).toBe(true);
  });

  it("应该检测未修改的 header", () => {
    const headers = new Headers([["user-agent", "same"]]);
    const session = createSessionForHeaders(headers);

    expect(session.isHeaderModified("user-agent")).toBe(false);
  });

  it("应该处理不存在的 header", () => {
    const headers = new Headers();
    const session = createSessionForHeaders(headers);

    expect(session.isHeaderModified("x-custom")).toBe(false);
  });

  it("应该检测到被删除的 header", () => {
    const headers = new Headers([["user-agent", "original"]]);
    const session = createSessionForHeaders(headers);

    session.headers.delete("user-agent");

    expect(session.isHeaderModified("user-agent")).toBe(true);
  });

  it("应该检测到新增的 header", () => {
    const headers = new Headers();
    const session = createSessionForHeaders(headers);

    session.headers.set("x-new-header", "new-value");

    expect(session.isHeaderModified("x-new-header")).toBe(true);
  });

  it("应该区分空字符串和 null", () => {
    const headers = new Headers([["x-test", ""]]);
    const session = createSessionForHeaders(headers);

    session.headers.delete("x-test");

    expect(session.isHeaderModified("x-test")).toBe(true); // "" -> null
    expect(session.headers.get("x-test")).toBeNull();
  });
});

describe("ProxySession.isWarmupRequest", () => {
  it("应识别合法的 Warmup 请求（忽略大小写与首尾空格）", () => {
    const session = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestMessage: {
        model: "claude-sonnet-4-5-20250929",
        messages: [
          {
            role: "user",
            content: [
              {
                type: "text",
                text: "  WaRmUp  ",
                cache_control: { type: "ephemeral" },
              },
            ],
          },
        ],
      },
    });

    expect(session.isWarmupRequest()).toBe(true);
  });

  it("endpoint 非 /v1/messages 时不应命中", () => {
    const session = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestUrl: new URL("http://localhost/v1/messages/count_tokens"),
      requestMessage: {
        messages: [
          {
            role: "user",
            content: [
              {
                type: "text",
                text: "Warmup",
                cache_control: { type: "ephemeral" },
              },
            ],
          },
        ],
      },
    });

    expect(session.isWarmupRequest()).toBe(false);
  });

  it("缺少 cache_control 或 type 不为 ephemeral 时不应命中", () => {
    const missingCacheControl = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestMessage: {
        messages: [
          {
            role: "user",
            content: [{ type: "text", text: "Warmup" }],
          },
        ],
      },
    });
    expect(missingCacheControl.isWarmupRequest()).toBe(false);

    const wrongCacheControl = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestMessage: {
        messages: [
          {
            role: "user",
            content: [
              {
                type: "text",
                text: "Warmup",
                cache_control: { type: "persistent" },
              },
            ],
          },
        ],
      },
    });
    expect(wrongCacheControl.isWarmupRequest()).toBe(false);
  });

  it("messages/content 非严格形态时不应命中（防误判）", () => {
    const multiMessages = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestMessage: {
        messages: [
          {
            role: "user",
            content: [
              {
                type: "text",
                text: "Warmup",
                cache_control: { type: "ephemeral" },
              },
            ],
          },
          {
            role: "user",
            content: [
              {
                type: "text",
                text: "Warmup",
                cache_control: { type: "ephemeral" },
              },
            ],
          },
        ],
      },
    });
    expect(multiMessages.isWarmupRequest()).toBe(false);

    const multiBlocks = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestMessage: {
        messages: [
          {
            role: "user",
            content: [
              {
                type: "text",
                text: "Warmup",
                cache_control: { type: "ephemeral" },
              },
              { type: "text", text: "Warmup", cache_control: { type: "ephemeral" } },
            ],
          },
        ],
      },
    });
    expect(multiBlocks.isWarmupRequest()).toBe(false);
  });

  it("messages/role/content 结构异常时不应命中", () => {
    const missingMessages = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestMessage: {},
    });
    expect(missingMessages.isWarmupRequest()).toBe(false);

    const nonArrayMessages = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestMessage: { messages: "Warmup" },
    });
    expect(nonArrayMessages.isWarmupRequest()).toBe(false);

    const roleNotUser = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestMessage: {
        messages: [
          {
            role: "assistant",
            content: [
              {
                type: "text",
                text: "Warmup",
                cache_control: { type: "ephemeral" },
              },
            ],
          },
        ],
      },
    });
    expect(roleNotUser.isWarmupRequest()).toBe(false);

    const contentNotArray = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestMessage: {
        messages: [{ role: "user", content: "Warmup" }],
      },
    });
    expect(contentNotArray.isWarmupRequest()).toBe(false);
  });

  it("block/text/cache_control 结构异常时不应命中", () => {
    const blockNotObject = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestMessage: {
        messages: [{ role: "user", content: [null] }],
      },
    });
    expect(blockNotObject.isWarmupRequest()).toBe(false);

    const typeNotText = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestMessage: {
        messages: [
          {
            role: "user",
            content: [
              {
                type: "image",
                text: "Warmup",
                cache_control: { type: "ephemeral" },
              },
            ],
          },
        ],
      },
    });
    expect(typeNotText.isWarmupRequest()).toBe(false);

    const textNotString = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestMessage: {
        messages: [
          {
            role: "user",
            content: [
              {
                type: "text",
                text: 123,
                cache_control: { type: "ephemeral" },
              },
            ],
          },
        ],
      },
    });
    expect(textNotString.isWarmupRequest()).toBe(false);

    const cacheControlNotObject = createSession({
      redirectedModel: "claude-sonnet-4-5-20250929",
      requestMessage: {
        messages: [
          {
            role: "user",
            content: [
              {
                type: "text",
                text: "Warmup",
                cache_control: null,
              },
            ],
          },
        ],
      },
    });
    expect(cacheControlNotObject.isWarmupRequest()).toBe(false);
  });
});

describe("ProxySession.addProviderToChain - endpoint audit", () => {
  it("应写入 vendorId/providerType/endpointId/endpointUrl", () => {
    const session = createSession({ redirectedModel: null });
    const provider = {
      id: 1,
      name: "p1",
      providerVendorId: 123,
      providerType: "claude",
      priority: 0,
      weight: 1,
      costMultiplier: 1,
      groupTag: null,
    } as unknown as Provider;

    session.addProviderToChain(provider, {
      endpointId: 42,
      endpointUrl: "https://api.example.com",
    });

    const chain = session.getProviderChain();
    expect(chain).toHaveLength(1);
    expect(chain[0]).toEqual(
      expect.objectContaining({
        id: 1,
        name: "p1",
        vendorId: 123,
        providerType: "claude",
        endpointId: 42,
        endpointUrl: "https://api.example.com",
      })
    );
  });

  it("同一 provider 连续写入且不带 attemptNumber 时应去重", () => {
    const session = createSession({ redirectedModel: null });
    const provider = {
      id: 1,
      name: "p1",
      providerVendorId: 123,
      providerType: "claude",
      priority: 0,
      weight: 1,
      costMultiplier: 1,
      groupTag: null,
    } as unknown as Provider;

    session.addProviderToChain(provider, { endpointId: 1, endpointUrl: "https://a.example.com" });
    session.addProviderToChain(provider, { endpointId: 2, endpointUrl: "https://b.example.com" });

    const chain = session.getProviderChain();
    expect(chain).toHaveLength(1);
    expect(chain[0]).toEqual(
      expect.objectContaining({
        endpointId: 1,
        endpointUrl: "https://a.example.com",
      })
    );
  });

  it("同一 provider 连续写入且带 attemptNumber 时应保留多条", () => {
    const session = createSession({ redirectedModel: null });
    const provider = {
      id: 1,
      name: "p1",
      providerVendorId: 123,
      providerType: "claude",
      priority: 0,
      weight: 1,
      costMultiplier: 1,
      groupTag: null,
    } as unknown as Provider;

    session.addProviderToChain(provider, {
      attemptNumber: 1,
      endpointId: 1,
      endpointUrl: "https://a.example.com",
    });
    session.addProviderToChain(provider, {
      attemptNumber: 2,
      endpointId: 2,
      endpointUrl: "https://b.example.com",
    });

    const chain = session.getProviderChain();
    expect(chain).toHaveLength(2);
    expect(chain[0]).toEqual(
      expect.objectContaining({
        attemptNumber: 1,
        endpointId: 1,
        endpointUrl: "https://a.example.com",
      })
    );
    expect(chain[1]).toEqual(
      expect.objectContaining({
        attemptNumber: 2,
        endpointId: 2,
        endpointUrl: "https://b.example.com",
      })
    );
  });
});