UHD_Media_Tool/packages/agent/test/upstream.test.ts

import { PI_BUILTIN_PROVIDER_IDS, PiProviderId } from "@umt/contracts";
import { describe, expect, it } from "vitest";
import {
  FakeOAuthControl,
  PiUpstreamClient,
  providerCatalog,
  resolveAgentSettings,
  type UpstreamAuthEvent,
  type UpstreamCredential,
  type UpstreamCredentialStore,
  UpstreamError,
  type UpstreamLoginInteraction,
  type UpstreamSpec,
} from "../src/index.ts";

/** 内存凭据存储:promise 链串行化 modify(模拟行锁) */
function memStore(initial?: UpstreamCredential) {
  let cur = initial;
  let chain: Promise<unknown> = Promise.resolve();
  const store: UpstreamCredentialStore & { get(): UpstreamCredential | undefined; writes: number } = {
    writes: 0,
    get: () => cur,
    read: async () => cur,
    modify: (fn) => {
      const run = chain.then(async () => {
        const next = await fn(cur);
        if (next !== undefined) {
          cur = next;
          store.writes += 1;
        }
        return cur;
      });
      chain = run.catch(() => undefined);
      return run;
    },
    delete: async () => {
      cur = undefined;
    },
  };
  return store;
}

const spec = (s: Partial<UpstreamSpec> & Pick<UpstreamSpec, "provider" | "authKind">): UpstreamSpec => ({
  id: "u1",
  ...s,
});

describe("能力表", () => {
  it("42 个内置 provider + openai-compatible + fake,id 与合同枚举一致", () => {
    const c = providerCatalog({ models: false });
    const ids = c.providers.map((p) => p.id);
    expect(ids).toHaveLength(44);
    expect(new Set(ids)).toEqual(new Set(PiProviderId.options));
    expect(ids.filter((i) => i !== "fake" && i !== "openai-compatible").sort()).toEqual(
      [...PI_BUILTIN_PROVIDER_IDS].sort(),
    );
    expect(c.catalogGeneratedAt).toMatch(/^\d{4}-/);
    expect(c.providers.every((p) => p.models.length === 0)).toBe(true);
  });

  it("订阅登录:9 种 OAuth;Codex 只有 oauth 且优先 device;OpenAI 支持原生 ChatGPT 登录", () => {
    const c = providerCatalog();
    const oauth = c.providers.filter((p) => p.oauth && p.id !== "fake").map((p) => p.id);
    expect(oauth.sort()).toEqual(
      [
        "anthropic",
        "github-copilot",
        "kimi-coding",
        "meta",
        "openai",
        "openai-codex",
        "openrouter",
        "radius",
        "xai",
      ].sort(),
    );
    const byId = Object.fromEntries(c.providers.map((p) => [p.id, p]));
    expect(byId["openai-codex"]!.authKinds).toEqual(["oauth"]);
    expect(byId["openai-codex"]!.oauth!.modes).toEqual(["device", "browser"]);
    expect(byId["openai-codex"]!.baseUrlOverridable).toBe(false);
    expect(byId.anthropic!.termsNote).toMatch(/Claude Code.*封禁/);
    expect(byId.anthropic!.group).toBe("subscription");
    for (const id of ["anthropic", "openai-codex", "github-copilot", "xai", "kimi-coding", "meta"])
      expect(byId[id]!.termsNote, id).toBeTruthy();
    expect(byId.openrouter!.termsNote).toBeNull();
    expect(byId.opencode!.oauth).toBeNull();
    expect(byId["opencode-go"]!.authKinds).toEqual(["api_key"]);
    expect(byId.radius!.dynamicModels).toBe(true);
    expect(byId["cloudflare-workers-ai"]!.envFields.map((f) => f.name)).toEqual(["CLOUDFLARE_ACCOUNT_ID"]);
    expect(byId["amazon-bedrock"]!.authKinds).toContain("none");
    expect(byId.openai!.models.length).toBeGreaterThan(10);
    expect(byId.openai!.baseUrlOverridable).toBe(true);
    expect(byId.openai!.authKinds).toEqual(["oauth", "api_key"]);
    expect(byId.openai!.oauth!.modes).toEqual(["browser"]);
    expect(byId.openai!.termsNote).toBeNull();
    expect(byId.typesafe!.models).toEqual([]);
  });
});

describe("按上游装配", () => {
  it("同一 provider 两个账号并存:各自的 key 只进各自的请求;不读宿主环境变量", async () => {
    process.env.OPENAI_API_KEY = "sk-host-env-should-not-leak";
    try {
      const s = resolveAgentSettings({ provider: "openai", model: "gpt-5", contextBudgetTokens: 64_000 });
      const a = new PiUpstreamClient(
        spec({ id: "a", provider: "openai", authKind: "api_key", model: "gpt-5" }),
        memStore({ type: "api_key", key: "sk-account-a" }),
      );
      const b = new PiUpstreamClient(
        spec({ id: "b", provider: "openai", authKind: "api_key", model: "gpt-5", baseUrl: "https://proxy.local/v1/" }),
        memStore({ type: "api_key", key: "sk-account-b" }),
      );
      const ra = await a.resolveRunModel(s);
      const rb = await b.resolveRunModel(s);
      expect(JSON.stringify(await ra.models.getAuth(ra.model))).toContain("sk-account-a");
      expect(JSON.stringify(await rb.models.getAuth(rb.model))).toContain("sk-account-b");
      expect(rb.model.baseUrl).toBe("https://proxy.local/v1");
      expect(ra.model.baseUrl).toBe("https://api.openai.com/v1");
      expect(ra.model.contextWindow).toBe(64_000);
      expect(JSON.stringify(ra.model)).not.toContain("sk-account");

      const empty = new PiUpstreamClient(spec({ provider: "openai", authKind: "api_key" }), memStore());
      await expect(empty.resolveAuth()).rejects.toMatchObject({ kind: "auth" });
    } finally {
      delete process.env.OPENAI_API_KEY;
    }
  });

  it("目录外的模型 id 以同 provider 模型为模板合成;思考档位夹到模型支持范围", async () => {
    const c = new PiUpstreamClient(
      spec({ provider: "anthropic", authKind: "api_key", model: "claude-future-9", thinkingLevel: "max" }),
      memStore({ type: "api_key", key: "sk-ant-x" }),
    );
    const r = await c.resolveRunModel(
      resolveAgentSettings({ provider: "anthropic", model: "ignored", thinkingLevel: "xhigh" }),
    );
    expect(r.model.id).toBe("claude-future-9");
    expect(r.thinkingLevel).toBeDefined();
    const codex = new PiUpstreamClient(
      spec({ provider: "openai-codex", authKind: "oauth", model: "gpt-5.1-codex-mini" }),
      memStore(),
    );
    const rc = await codex.resolveRunModel(resolveAgentSettings({ provider: "x", model: "m", thinkingLevel: "max" }));
    expect(rc.thinkingLevel).not.toBe("max");
    expect(r.model.api).toBe("anthropic-messages");
    const t = await c.test({ mode: "auth" });
    expect(t).toMatchObject({ model: "claude-future-9", inCatalog: false, source: "stored credential" });
  });

  it("测试补全带 sessionId:OpenCode Go 请求头有 x-opencode-session", async () => {
    const seen: { url: string; headers: Headers }[] = [];
    const orig = globalThis.fetch;
    globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => {
      const req = new Request(input, init);
      seen.push({ url: req.url, headers: req.headers });
      return new Response(JSON.stringify({ error: { message: "stubbed" } }), {
        status: 500,
        headers: { "content-type": "application/json" },
      });
    }) as typeof fetch;
    try {
      const c = new PiUpstreamClient(
        spec({ provider: "opencode-go", authKind: "api_key" }),
        memStore({ type: "api_key", key: "sk-opencode-test" }),
      );
      await expect(c.test({ mode: "completion" })).rejects.toBeInstanceOf(UpstreamError);
    } finally {
      globalThis.fetch = orig;
    }
    const call = seen.find((x) => x.headers.get("x-opencode-session"));
    expect(call, `请求:${seen.map((x) => x.url).join(", ")}`).toBeDefined();
    expect(call?.headers.get("x-opencode-session")).toMatch(/^umt-upstream-test-/);
  });

  it("自定义兼容端点:baseUrl/模型必填;authKind none 不需要 key", async () => {
    expect(() => new PiUpstreamClient(spec({ provider: "openai-compatible", authKind: "none" }), memStore())).toThrow(
      UpstreamError,
    );
    const c = new PiUpstreamClient(
      spec({ provider: "openai-compatible", authKind: "none", baseUrl: "http://llm.local/v1", model: "qwen3" }),
      memStore(),
    );
    expect((await c.resolveAuth()).source).toBe("none");
    const m = await c.listModels();
    expect(m.items).toMatchObject([{ id: "qwen3", source: "custom", api: "openai-completions" }]);
  });
});

/** 挂起式交互:按模式回答 select,manual_code 等外部粘贴 */
function interaction(mode: "browser" | "device_code") {
  const events: UpstreamAuthEvent[] = [];
  let paste: ((v: string) => void) | null = null;
  const first = Promise.withResolvers<UpstreamAuthEvent>();
  const it: UpstreamLoginInteraction = {
    prompt: async (p) => {
      if (p.type === "select") return mode;
      if (p.type === "manual_code") return await new Promise<string>((r) => (paste = r));
      throw new Error(`unexpected prompt ${p.type}`);
    },
    notify: (e) => {
      events.push(e);
      first.resolve(e);
    },
  };
  return { it, events, first: first.promise, paste: (v: string) => paste?.(v) };
}

describe("fake provider 的模拟 OAuth", () => {
  it("browser:粘回地址栏 URL(state 校验)→ 写入 token 包;device:批准后写入", async () => {
    const control = new FakeOAuthControl();
    control.deviceAutoApproveMs = null;
    const store = memStore();
    const c = new PiUpstreamClient(spec({ provider: "fake", authKind: "oauth" }), store, { fakeOAuth: control });
    await expect(c.resolveAuth()).rejects.toMatchObject({ kind: "auth" });

    const b = interaction("browser");
    const login = c.login(b.it);
    const e = await b.first;
    expect(e.type).toBe("auth_url");
    const state = new URL((e as { url: string }).url).searchParams.get("state");
    b.paste(`http://localhost:1455/auth/callback?code=abc&state=${state}`);
    await login;
    expect(store.get()).toMatchObject({ type: "oauth", account: "fake-browser-user" });
    expect((await c.resolveAuth()).source).toBe("OAuth");

    const bad = interaction("browser");
    const badLogin = c.login(bad.it);
    await bad.first;
    bad.paste("http://localhost:1455/auth/callback?code=abc&state=wrong");
    await expect(badLogin).rejects.toThrow(/state mismatch/);

    const d = interaction("device_code");
    const dl = c.login(d.it);
    const de = await d.first;
    expect(de).toMatchObject({ type: "device_code", verificationUri: "https://fake-oauth.invalid/device" });
    control.approveDevice();
    await dl;
    expect(store.get()).toMatchObject({ account: "fake-device-user" });
  });

  it("刷新:过期前锁内刷新,并发只刷新一次;refresh token 轮换后写回;失败保留旧凭据并报 oauth", async () => {
    const control = new FakeOAuthControl();
    control.refreshDelayMs = 30;
    const store = memStore({ type: "oauth", access: "old-a", refresh: "old-r", expires: Date.now() + 60_000 });
    const c1 = new PiUpstreamClient(spec({ provider: "fake", authKind: "oauth" }), store, { fakeOAuth: control });
    const c2 = new PiUpstreamClient(spec({ provider: "fake", authKind: "oauth" }), store, { fakeOAuth: control });
    await Promise.all([c1.resolveAuth(30 * 60_000), c2.resolveAuth(30 * 60_000), c1.resolveAuth(30 * 60_000)]);
    expect(control.refreshCalls).toBe(1);
    const cur = store.get() as { refresh: string; expires: number };
    expect(cur.refresh).not.toBe("old-r");
    expect(cur.expires).toBeGreaterThan(Date.now() + 30 * 60_000);

    control.failRefresh = true;
    const before = store.get();
    await expect(c1.resolveAuth(2 * 3600_000)).rejects.toMatchObject({ kind: "oauth" });
    expect(store.get()).toEqual(before);
  });
});