UHD_Media_Tool/apps/server/src/modules/pi/oauth.ts

import { randomUUID } from "node:crypto";
import {
  classifyUpstreamError,
  OAUTH_CALLBACK_PORTS,
  PiUpstreamClient,
  providerInfo,
  type UpstreamAuthEvent,
  type UpstreamAuthPrompt,
  type UpstreamClientOptions,
  type UpstreamLoginInteraction,
} from "@umt/agent";
import type { PiOAuthMode, PiOAuthStartResult, PiOAuthState } from "@umt/contracts";
import type { PiUpstreamRow } from "@umt/db";
import { ApiException } from "../../http/errors.ts";
import type { ServerContext } from "../types.ts";
import { dbCredentialStore, getRow, setStatus, specOf } from "./store.ts";

// 服务端订阅登录(研究文档 §3):调用 pi 的 models.login,注入“挂起式” interaction ——
// select 用 start 请求里的 mode 作答、text 用 enterpriseDomain 作答、manual_code 挂起等 /oauth/complete 粘贴;
// notify(auth_url | device_code) 写进会话记录,start 拿到第一个事件就返回。
// 会话只活在当前 api 进程内存(interaction 是闭包):多实例部署时登录请求要粘到同一实例。
// 同一 provider 的登录串行(pi 的浏览器流程会在服务端监听固定回调端口,并发会撞端口)。

export interface OAuthOptions {
  /** 会话有效期(默认 15 分钟) */
  sessionTtlMs?: number;
  /** start 等第一个事件的上限(默认 20s) */
  startWaitMs?: number;
  /** complete 缺省等待(默认 25s) */
  completeWaitMs?: number;
  client?: UpstreamClientOptions;
}

interface Session {
  id: string;
  upstreamId: string;
  provider: string;
  mode: PiOAuthMode;
  state: PiOAuthState;
  message: string | null;
  progress: string[];
  authorizeUrl: string | null;
  userCode: string | null;
  verificationUrl: string | null;
  instructions: string | null;
  expiresAt: number;
  abort: AbortController;
  /** manual_code 挂起的解析函数 */
  paste: ((v: string) => void) | null;
  pasteWaiters: (() => void)[];
  firstEvent: PromiseWithResolvers<void>;
  done: Promise<void>;
  timer: ReturnType<typeof setTimeout>;
}

const MODE_RE: Record<PiOAuthMode, RegExp> = { browser: /browser/i, device: /device/i };

/** 登录失败信息转成可读提示(端口占用、取消、超时) */
export function readableLoginError(provider: string, err: unknown): string {
  const raw = err instanceof Error ? err.message : String(err);
  if (/EADDRINUSE|address already in use/i.test(raw)) {
    const port = OAUTH_CALLBACK_PORTS[provider] ?? raw.match(/:(\d{2,5})\b/)?.[1];
    return `服务端回调端口${port ? ` ${port}` : ""} 被占用。结束占用该端口的进程（如另一个登录流程、Claude Code / Codex CLI）后重试`;
  }
  if (/cancel|abort/i.test(raw)) return "登录已取消";
  if (/state mismatch/i.test(raw)) return "粘贴的地址不属于本次登录（state 不一致）。用本次打开的授权链接重新登录";
  if (/\b404\b/.test(raw) && provider === "openai-codex")
    return "这个账号未开启设备码登录（Codex 返回 404）。改用浏览器登录";
  return classifyUpstreamError(err).message;
}

export class OAuthSessions {
  private readonly sessions = new Map<string, Session>();
  constructor(
    private readonly ctx: ServerContext,
    private readonly opts: OAuthOptions = {},
  ) {}

  private active(provider: string): Session | undefined {
    for (const s of this.sessions.values()) if (s.provider === provider && s.state === "pending") return s;
    return undefined;
  }

  private finish(s: Session, state: PiOAuthState, message: string | null) {
    if (s.state !== "pending") return;
    s.state = state;
    s.message = message;
    s.firstEvent.resolve();
    for (const w of s.pasteWaiters.splice(0)) w();
    clearTimeout(s.timer);
    // 结束后保留一会儿供 complete 查询结果
    s.timer = setTimeout(() => this.sessions.delete(s.id), 5 * 60_000);
    (s.timer as { unref?: () => void }).unref?.();
  }

  cancel(s: Session, reason: PiOAuthState = "cancelled", message = "登录已取消") {
    this.finish(s, reason, message);
    s.abort.abort(new Error(message));
  }

  cancelUpstream(upstreamId: string) {
    for (const s of this.sessions.values()) if (s.upstreamId === upstreamId && s.state === "pending") this.cancel(s);
  }

  async start(row: PiUpstreamRow, req: { mode?: PiOAuthMode; enterpriseDomain?: string }): Promise<PiOAuthStartResult> {
    const info = providerInfo(row.provider);
    if (row.auth_kind !== "oauth" || !info?.oauth)
      throw new ApiException("precondition_failed", "这个模型服务不是订阅登录类型");
    if (!row.enabled) throw new ApiException("precondition_failed", "模型服务已停用");
    const mode = req.mode ?? info.oauth.modes[0]!;
    if (!info.oauth.modes.includes(mode))
      throw new ApiException(
        "validation_failed",
        `${info.name} 不支持 ${mode} 登录，可用：${info.oauth.modes.join(" / ")}`,
      );

    const busy = this.active(row.provider);
    if (busy) {
      if (busy.upstreamId !== row.id)
        throw new ApiException(
          "conflict",
          `${info.name} 已有登录进行中。完成它，或等到 ${new Date(busy.expiresAt).toISOString()} 过期后重试`,
        );
      this.cancel(busy, "cancelled", "已被新的登录请求取代");
      await busy.done.catch(() => undefined);
    }

    const ttl = this.opts.sessionTtlMs ?? 15 * 60_000;
    const s: Session = {
      id: randomUUID(),
      upstreamId: row.id,
      provider: row.provider,
      mode,
      state: "pending",
      message: null,
      progress: [],
      authorizeUrl: null,
      userCode: null,
      verificationUrl: null,
      instructions: null,
      expiresAt: Date.now() + ttl,
      abort: new AbortController(),
      paste: null,
      pasteWaiters: [],
      firstEvent: Promise.withResolvers<void>(),
      done: Promise.resolve(),
      timer: setTimeout(() => this.cancel(s, "expired", "登录会话已超时，重新开始登录"), ttl),
    };
    (s.timer as { unref?: () => void }).unref?.();
    this.sessions.set(s.id, s);

    const interaction: UpstreamLoginInteraction = {
      signal: s.abort.signal,
      prompt: (p: UpstreamAuthPrompt) => this.prompt(s, p, req),
      notify: (e: UpstreamAuthEvent) => this.notify(s, e),
    };
    const client = new PiUpstreamClient(specOf(row), dbCredentialStore(this.ctx, row.id), this.opts.client);
    s.done = client.login(interaction).then(
      async () => {
        await setStatus(this.ctx, row.id, { status: "ok", verified: true });
        this.finish(s, "succeeded", null);
      },
      async (err) => {
        if (s.state !== "pending") return; // 已取消/超时
        const message = readableLoginError(row.provider, err);
        this.ctx.logger.warn({ upstreamId: row.id, provider: row.provider, message }, "订阅登录失败");
        this.finish(s, "failed", message);
      },
    );

    const waitMs = this.opts.startWaitMs ?? 20_000;
    await Promise.race([s.firstEvent.promise, new Promise((r) => setTimeout(r, waitMs).unref?.())]);
    return this.startView(s);
  }

  private async prompt(s: Session, p: UpstreamAuthPrompt, req: { mode?: PiOAuthMode; enterpriseDomain?: string }) {
    if (p.type === "select") {
      const hit = p.options.find((o) => MODE_RE[s.mode].test(`${o.id} ${o.label}`));
      if (!hit) throw new Error(`provider 不提供 ${s.mode} 登录方式`);
      return hit.id;
    }
    if (p.type === "text") return req.enterpriseDomain ?? "";
    if (p.type === "manual_code") {
      return await new Promise<string>((resolve, reject) => {
        const onAbort = () => reject(new Error("Login cancelled"));
        p.signal?.addEventListener("abort", onAbort, { once: true });
        s.abort.signal.addEventListener("abort", onAbort, { once: true });
        s.paste = (v) => {
          s.paste = null;
          resolve(v);
        };
        for (const w of s.pasteWaiters.splice(0)) w();
      });
    }
    throw new Error("这个登录流程需要额外输入(secret)，服务端不支持");
  }

  private notify(s: Session, e: UpstreamAuthEvent) {
    if (e.type === "auth_url") {
      s.authorizeUrl = e.url;
      s.instructions = e.instructions ?? null;
      s.mode = "browser";
      s.firstEvent.resolve();
    } else if (e.type === "device_code") {
      s.userCode = e.userCode;
      s.verificationUrl = e.verificationUri;
      s.mode = "device";
      if (e.expiresInSeconds) s.expiresAt = Math.min(s.expiresAt, Date.now() + e.expiresInSeconds * 1000);
      s.firstEvent.resolve();
    } else {
      s.progress.push(e.message.slice(0, 300));
      if (s.progress.length > 20) s.progress.shift();
      if (e.type === "info" && !s.instructions) s.instructions = e.message.slice(0, 500);
    }
  }

  private startView(s: Session): PiOAuthStartResult {
    return {
      sessionId: s.id,
      mode: s.mode,
      state: s.state,
      authorizeUrl: s.authorizeUrl,
      userCode: s.userCode,
      verificationUrl: s.verificationUrl,
      instructions: s.instructions,
      expiresAt: new Date(s.expiresAt).toISOString(),
      message: s.message,
    };
  }

  /** browser:把粘贴内容交给挂起的 manual_code;然后(两种模式)等结果到 waitMs,仍未完成返回 pending */
  async complete(
    upstreamId: string,
    req: { sessionId: string; pastedUrlOrCode?: string; waitMs?: number },
  ): Promise<{ sessionId: string; state: PiOAuthState; message: string | null; progress: string[] }> {
    const s = this.sessions.get(req.sessionId);
    if (!s || s.upstreamId !== upstreamId)
      return {
        sessionId: req.sessionId,
        state: "expired",
        message: "登录会话不存在或已过期，重新开始登录",
        progress: [],
      };
    if (req.pastedUrlOrCode && s.state === "pending") {
      if (s.mode !== "browser") throw new ApiException("bad_request", "设备码登录不需要粘贴，等待授权完成");
      if (!s.paste) {
        // manual_code 提示可能还没发出:短暂等待
        await Promise.race([
          new Promise<void>((r) => s.pasteWaiters.push(r)),
          new Promise((r) => setTimeout(r, 2000).unref?.()),
        ]);
      }
      if (!s.paste) throw new ApiException("conflict", "当前登录不在等待粘贴的阶段");
      s.paste(req.pastedUrlOrCode);
    }
    const waitMs = req.waitMs ?? this.opts.completeWaitMs ?? 25_000;
    if (s.state === "pending")
      await Promise.race([s.done, new Promise((r) => setTimeout(r, waitMs).unref?.())]).catch(() => undefined);
    return { sessionId: s.id, state: s.state, message: s.message, progress: [...s.progress] };
  }

  async shutdown() {
    for (const s of this.sessions.values()) {
      if (s.state === "pending") this.cancel(s, "cancelled", "服务正在停止");
      clearTimeout(s.timer);
    }
    this.sessions.clear();
  }
}

export async function requireRow(ctx: ServerContext, id: string): Promise<PiUpstreamRow> {
  const row = await getRow(ctx, id);
  if (!row) throw new ApiException("not_found", "模型服务不存在");
  return row;
}