UHD_Media_Tool/packages/agent/src/pi/fake-upstream.ts

import { randomBytes } from "node:crypto";
import {
  type AssistantMessage,
  createProvider,
  type FauxResponseFactory,
  fauxAssistantMessage,
  fauxProvider,
  type OAuthAuth,
  type OAuthCredential,
  type Provider,
  type ProviderAuthInteraction,
} from "@earendil-works/pi-ai";

// provider = "fake" 的上游:faux 模型 + 模拟 OAuth(browser 粘 URL / device 轮询),不联网。
// 设置页联调与 OAuth 状态机测试用;行为由 FakeOAuthControl 控制(测试注入,缺省用进程级默认实例)。

export class FakeOAuthControl {
  loginCalls = 0;
  refreshCalls = 0;
  /** refresh 抛 invalid_grant */
  failRefresh = false;
  /** refresh 前等待(并发测试用) */
  refreshDelayMs = 0;
  /** 新 token 的有效期 */
  tokenTtlMs = 3600_000;
  /** device 模式:多久后自动批准(null = 只能 approveDevice() 批准) */
  deviceAutoApproveMs: number | null = 3000;
  /** browser 模式:模拟回调端口被占用(pi 的 Anthropic 会直接 reject) */
  browserListenError: string | null = null;
  private approved = false;

  approveDevice(): void {
    this.approved = true;
  }

  /** @internal */
  consumeApproval(): boolean {
    if (!this.approved) return false;
    this.approved = false;
    return true;
  }
}

export const DEFAULT_FAKE_OAUTH = new FakeOAuthControl();

const sleep = (ms: number, signal?: AbortSignal) =>
  new Promise<void>((resolve, reject) => {
    const t = setTimeout(resolve, ms);
    signal?.addEventListener(
      "abort",
      () => {
        clearTimeout(t);
        reject(new Error("Login cancelled"));
      },
      { once: true },
    );
  });

function parsePasted(input: string): { code?: string; state?: string } {
  const v = input.trim();
  try {
    const u = new URL(v);
    return { code: u.searchParams.get("code") ?? undefined, state: u.searchParams.get("state") ?? undefined };
  } catch {
    // code#state 或裸 code
  }
  if (v.includes("#")) {
    const [code, state] = v.split("#", 2);
    return { code: code || undefined, state: state || undefined };
  }
  if (v.includes("code=")) {
    const p = new URLSearchParams(v);
    return { code: p.get("code") ?? undefined, state: p.get("state") ?? undefined };
  }
  return { code: v || undefined };
}

function issue(control: FakeOAuthControl, account: string): OAuthCredential {
  const n = randomBytes(6).toString("hex");
  return {
    type: "oauth",
    access: `fake-access-${n}`,
    refresh: `fake-refresh-${n}`,
    expires: Date.now() + control.tokenTtlMs,
    account,
  };
}

export function fakeOAuth(control: FakeOAuthControl): OAuthAuth {
  return {
    name: "Fake OAuth",
    isSubscription: false,
    async login(interaction: ProviderAuthInteraction) {
      control.loginCalls += 1;
      const method = await interaction.prompt({
        type: "select",
        message: "Select fake login method:",
        options: [
          { id: "browser", label: "Browser login" },
          { id: "device_code", label: "Device code login" },
        ],
      });
      if (method === "device_code") {
        const userCode = `FAKE-${randomBytes(2).toString("hex").toUpperCase()}`;
        interaction.notify({
          type: "device_code",
          userCode,
          verificationUri: "https://fake-oauth.invalid/device",
          intervalSeconds: 1,
          expiresInSeconds: 900,
        });
        const started = Date.now();
        for (;;) {
          if (control.consumeApproval()) break;
          if (control.deviceAutoApproveMs !== null && Date.now() - started >= control.deviceAutoApproveMs) break;
          await sleep(25, interaction.signal);
        }
        interaction.notify({ type: "progress", message: "Device authorized" });
        return issue(control, "fake-device-user");
      }
      if (control.browserListenError) throw new Error(control.browserListenError);
      const state = randomBytes(8).toString("hex");
      interaction.notify({
        type: "auth_url",
        url: `https://fake-oauth.invalid/authorize?state=${state}&redirect_uri=${encodeURIComponent("http://localhost:1455/auth/callback")}`,
        instructions: "Complete login in your browser, then paste the final redirect URL here.",
      });
      const input = await interaction.prompt({
        type: "manual_code",
        message: "Paste the authorization code / redirect URL:",
        placeholder: "http://localhost:1455/auth/callback",
      });
      const parsed = parsePasted(input);
      if (!parsed.code) throw new Error("Missing authorization code");
      if (parsed.state && parsed.state !== state) throw new Error("OAuth state mismatch");
      interaction.notify({ type: "progress", message: "Exchanging authorization code for tokens..." });
      return issue(control, "fake-browser-user");
    },
    async refresh(credential, signal) {
      control.refreshCalls += 1;
      if (control.refreshDelayMs > 0) await sleep(control.refreshDelayMs, signal);
      if (control.failRefresh) throw new Error("invalid_grant: refresh token revoked");
      return issue(control, String(credential.account ?? "fake-user"));
    },
    async toAuth(credential) {
      return { apiKey: credential.access };
    },
  };
}

/** provider id "fake" 的上游 provider:faux 流 + 始终就绪的 apiKey(authKind none)+ 模拟 OAuth */
export function fakeUpstreamProvider(control: FakeOAuthControl): Provider {
  const faux = fauxProvider({
    provider: "fake",
    models: [{ id: "fake-model", contextWindow: 64_000, maxTokens: 4096, reasoning: true }],
    tokensPerSecond: 0,
  });
  const factory: FauxResponseFactory = (): AssistantMessage => {
    faux.appendResponses([factory]);
    return fauxAssistantMessage("ok");
  };
  faux.setResponses([factory]);
  const base = faux.provider;
  return createProvider({
    id: "fake",
    name: "Fake",
    auth: { apiKey: { name: "Fake", resolve: async () => ({ auth: {}, source: "fake" }) }, oauth: fakeOAuth(control) },
    models: base.getModels(),
    api: { stream: base.stream.bind(base), streamSimple: base.streamSimple.bind(base) } as never,
  });
}