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

import { existsSync } from "node:fs";
import { homedir } from "node:os";
import process from "node:process";
import {
  type Api,
  type AuthContext,
  type AuthInteraction,
  type Credential,
  type CredentialStore,
  clampThinkingLevel,
  createModels,
  createProvider,
  type Model,
  ModelsError,
  type MutableModels,
  type Provider,
  type ProviderStreams,
} from "@earendil-works/pi-ai";
import { anthropicMessagesApi } from "@earendil-works/pi-ai/api/anthropic-messages.lazy";
import { googleGenerativeAIApi } from "@earendil-works/pi-ai/api/google-generative-ai.lazy";
import { openAICompletionsApi } from "@earendil-works/pi-ai/api/openai-completions.lazy";
import { openAIResponsesApi } from "@earendil-works/pi-ai/api/openai-responses.lazy";
import type { PiModelInfo, PiTestErrorKind } from "@umt/contracts";
import { redactString } from "../redact.ts";
import type { AgentThinkingLevel, ResolvedAgentSettings } from "../settings.ts";
import { builtinProvider, modelInfo, providerInfo } from "./catalog.ts";
import { DEFAULT_FAKE_OAUTH, type FakeOAuthControl, fakeUpstreamProvider } from "./fake-upstream.ts";
import type { ResolvedPiModel } from "./model.ts";

// 按上游实例装配 pi-ai:每个上游一个独立的 Models,里面只有它那一个 provider,
// credentials 是只映射到本上游 secret 的 CredentialStore 适配器 —— 绕开“pi 按 provider id 只存一份凭据”,
// 同一 provider 的多个账号可以并存。明文 key / token 只经 read() 交给 SDK,不进事件、日志、错误信息。

/** 上游描述(来自 pi_upstream 行,非密) */
export interface UpstreamSpec {
  id: string;
  provider: string;
  authKind: "api_key" | "oauth" | "none";
  baseUrl?: string | null;
  /** 缺省模型 id;null 用目录第一个 */
  model?: string | null;
  thinkingLevel?: string;
  extra?: {
    env?: Record<string, string>;
    api?: string;
    compat?: Record<string, unknown>;
    headers?: Record<string, string>;
    contextWindow?: number;
  };
}

/** 与 pi Credential 同形(本包自有类型,不暴露 SDK) */
export type UpstreamCredential =
  | { type: "api_key"; key?: string; env?: Record<string, string> }
  | { type: "oauth"; access: string; refresh: string; expires: number; [k: string]: unknown };

/**
 * 本上游凭据的存取。modify 是唯一写路径,实现必须跨进程互斥(server 用 pi_upstream 行锁),
 * 因为 refresh token 会轮换,并发刷新会让其中一方失效。fn 返回 undefined 表示不改。
 */
export interface UpstreamCredentialStore {
  read(): Promise<UpstreamCredential | undefined>;
  modify(
    fn: (current: UpstreamCredential | undefined) => Promise<UpstreamCredential | undefined>,
  ): Promise<UpstreamCredential | undefined>;
  delete(): Promise<void>;
}

export interface UpstreamClientOptions {
  /** provider = fake 的模拟 OAuth 行为(测试注入) */
  fakeOAuth?: FakeOAuthControl;
}

/** 登录交互(与 pi AuthInteraction 同形) */
export type UpstreamAuthPrompt = { signal?: AbortSignal } & (
  | { type: "text"; message: string; placeholder?: string }
  | { type: "secret"; message: string; placeholder?: string }
  | { type: "select"; message: string; options: readonly { id: string; label: string; description?: string }[] }
  | { type: "manual_code"; message: string; placeholder?: string }
);
export type UpstreamAuthEvent =
  | { type: "info"; message: string; links?: readonly { url: string; label?: string }[] }
  | { type: "auth_url"; url: string; instructions?: string }
  | {
      type: "device_code";
      userCode: string;
      verificationUri: string;
      intervalSeconds?: number;
      expiresInSeconds?: number;
    }
  | { type: "progress"; message: string };
export interface UpstreamLoginInteraction {
  signal?: AbortSignal;
  prompt(prompt: UpstreamAuthPrompt): Promise<string>;
  notify(event: UpstreamAuthEvent): void;
}

/** 上游错误(已脱敏);kind 供 UI 分类与状态机使用 */
export class UpstreamError extends Error {
  constructor(
    readonly kind: PiTestErrorKind,
    message: string,
  ) {
    super(message);
    this.name = "UpstreamError";
  }
}

export function classifyUpstreamError(err: unknown): UpstreamError {
  if (err instanceof UpstreamError) return err;
  const raw = err instanceof Error ? err.message : String(err);
  const message = redactString(raw).slice(0, 500);
  if (err instanceof ModelsError) {
    if (err.code === "oauth") return new UpstreamError("oauth", message);
    if (err.code === "auth") return new UpstreamError("auth", message);
    if (err.code === "provider" || err.code === "model_validation") return new UpstreamError("config", message);
  }
  if (/\b(401|403)\b|unauthori[sz]ed|invalid[_ ]api[_ ]key|authentication|forbidden/i.test(raw))
    return new UpstreamError("auth", message);
  if (/\b404\b|model[^\n]{0,40}(not found|does not exist|unknown)/i.test(raw))
    return new UpstreamError("model", message);
  if (
    /ECONNREFUSED|ENOTFOUND|EAI_AGAIN|ETIMEDOUT|ECONNRESET|fetch failed|connection error|network|timed? ?out|socket/i.test(
      raw,
    )
  )
    return new UpstreamError("network", message);
  return new UpstreamError("other", message);
}

const CUSTOM_APIS: Record<string, () => ProviderStreams> = {
  "openai-completions": openAICompletionsApi,
  "openai-responses": openAIResponsesApi,
  "anthropic-messages": anthropicMessagesApi,
  "google-generative-ai": googleGenerativeAIApi,
};

const trimSlash = (u: string) => u.replace(/\/+$/, "");

/** baseUrl 覆盖:Anthropic 形状去掉 /v1(SDK 自己加),其它按原样 */
function overrideBase(m: Model<Api>, baseUrl: string): Model<Api> {
  const b = trimSlash(baseUrl);
  return { ...m, baseUrl: m.api === "anthropic-messages" ? b.replace(/\/v1$/, "") : b };
}

function withBaseUrl(base: Provider, baseUrl: string): Provider {
  return {
    ...base,
    baseUrl,
    getModels: () => base.getModels().map((m) => overrideBase(m, baseUrl)),
  };
}

function customProvider(spec: UpstreamSpec): Provider {
  const api = spec.extra?.api ?? "openai-completions";
  const factory = CUSTOM_APIS[api];
  if (!factory) throw new UpstreamError("config", `不支持的 API 形状:${api}`);
  if (!spec.baseUrl) throw new UpstreamError("config", "自定义兼容端点必须配置 baseUrl");
  if (!spec.model) throw new UpstreamError("config", "自定义兼容端点必须配置模型 id");
  const base = api === "anthropic-messages" ? trimSlash(spec.baseUrl).replace(/\/v1$/, "") : trimSlash(spec.baseUrl);
  const model = {
    id: spec.model,
    name: spec.model,
    api,
    provider: "openai-compatible",
    baseUrl: base,
    reasoning: (spec.thinkingLevel ?? "off") !== "off",
    input: ["text"],
    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
    contextWindow: spec.extra?.contextWindow ?? 128_000,
    maxTokens: 8192,
    ...(spec.extra?.headers ? { headers: spec.extra.headers } : {}),
    ...(api === "openai-completions"
      ? { compat: { supportsDeveloperRole: false, ...(spec.extra?.compat ?? {}) } }
      : spec.extra?.compat
        ? { compat: spec.extra.compat }
        : {}),
  } as Model<Api>;
  const keyless = spec.authKind === "none";
  return createProvider({
    id: "openai-compatible",
    name: "自定义兼容端点",
    auth: {
      apiKey: {
        name: "API key",
        resolve: async ({ credential }) => {
          if (credential?.key) return { auth: { apiKey: credential.key }, source: "stored credential" };
          return keyless ? { auth: {}, source: "none" } : undefined;
        },
      },
    },
    models: [model],
    api: factory(),
  });
}

function providerFor(spec: UpstreamSpec, opts: UpstreamClientOptions): Provider {
  if (spec.provider === "fake") return fakeUpstreamProvider(opts.fakeOAuth ?? DEFAULT_FAKE_OAUTH);
  if (spec.provider === "openai-compatible") return customProvider(spec);
  const info = providerInfo(spec.provider);
  // 动态 provider(Radius)持有模型表状态,每个上游用新实例
  const base = builtinProvider(spec.provider, info?.dynamicModels === true);
  if (!base) throw new UpstreamError("config", `未知 provider:${spec.provider}`);
  if (spec.baseUrl && info?.baseUrlOverridable && spec.authKind !== "oauth") return withBaseUrl(base, spec.baseUrl);
  return base;
}

const expandHome = (p: string) => (p.startsWith("~") ? `${homedir()}${p.slice(1)}` : p);

/**
 * 一个上游的 pi 客户端:凭据解析(OAuth 按需刷新,锁内双重检查)、登录、登出、模型表、测试、运行时模型装配。
 */
export class PiUpstreamClient {
  readonly providerId: string;
  private readonly models: MutableModels;
  private readonly provider: Provider;

  constructor(
    readonly spec: UpstreamSpec,
    private readonly store: UpstreamCredentialStore,
    opts: UpstreamClientOptions = {},
  ) {
    this.provider = providerFor(spec, opts);
    this.providerId = this.provider.id;
    const id = this.providerId;
    const envOf = (): Record<string, string> => spec.extra?.env ?? {};
    const credentials: CredentialStore = {
      read: async (pid) => (pid === id ? this.toPi(await store.read()) : undefined),
      list: async () => {
        const c = await store.read();
        return c ? [{ providerId: id, type: c.type }] : [];
      },
      modify: async (pid, fn) => {
        if (pid !== id) throw new Error(`凭据存储只服务 provider ${id}`);
        const out = await store.modify(async (cur) => {
          const next = await fn(this.toPi(cur));
          return next as UpstreamCredential | undefined;
        });
        return this.toPi(out);
      },
      delete: async (pid) => {
        if (pid === id) await store.delete();
      },
    };
    // 只有 authKind none(云环境凭据链 / ADC)才读进程环境变量;其它只认本上游 extra.env,避免误用宿主的 key
    const ambient = spec.authKind === "none";
    const authContext: AuthContext = {
      env: async (name) => envOf()[name] || (ambient ? process.env[name] || undefined : undefined),
      fileExists: async (path) => (ambient ? existsSync(expandHome(path)) : false),
    };
    this.models = createModels({ credentials, authContext });
    this.models.setProvider(this.provider);
  }

  private toPi(c: UpstreamCredential | undefined): Credential | undefined {
    if (!c) return undefined;
    if (c.type === "api_key") {
      const env = { ...(this.spec.extra?.env ?? {}), ...(c.env ?? {}) };
      return { type: "api_key", ...(c.key ? { key: c.key } : {}), ...(Object.keys(env).length ? { env } : {}) };
    }
    return c as Credential;
  }

  /** 解析请求凭据;OAuth 剩余有效期不足 minValidityMs(至少 5 分钟)时在锁内刷新并写回 */
  async resolveAuth(minValidityMs?: number): Promise<{ source: string | null }> {
    if (this.spec.authKind === "oauth") {
      const c = await this.store.read();
      if (c?.type !== "oauth") throw new UpstreamError("auth", "尚未登录(需要订阅登录)");
    }
    let r: Awaited<ReturnType<MutableModels["getAuth"]>>;
    try {
      r = await this.models.getAuth(
        this.providerId,
        minValidityMs !== undefined ? { minOAuthValidityMs: minValidityMs } : undefined,
      );
    } catch (err) {
      throw classifyUpstreamError(err);
    }
    if (!r) throw new UpstreamError("auth", this.spec.authKind === "api_key" ? "未配置 API Key" : "凭据不完整");
    return { source: r.source ?? null };
  }

  /** 运行 provider 的订阅登录流程;成功后经 store.modify 写回 */
  async login(interaction: UpstreamLoginInteraction): Promise<void> {
    try {
      await this.models.login(this.providerId, "oauth", interaction as AuthInteraction);
    } catch (err) {
      throw classifyUpstreamError(err);
    }
  }

  async logout(): Promise<void> {
    await this.models.logout(this.providerId);
  }

  /** 模型表:动态 provider 先刷新(失败保留静态表) */
  async listModels(opts: { refresh?: boolean } = {}): Promise<{
    items: PiModelInfo[];
    refreshed: boolean;
    refreshError: string | null;
  }> {
    let refreshed = false;
    let refreshError: string | null = null;
    const dynamic = typeof this.provider.refreshModels === "function";
    if (dynamic && opts.refresh !== false) {
      try {
        const r = await this.models.refresh({ providers: [this.providerId], allowNetwork: true, force: true });
        const e = r.errors.get(this.providerId);
        if (e) refreshError = redactString(e.message).slice(0, 300);
        else refreshed = true;
      } catch (err) {
        refreshError = classifyUpstreamError(err).message;
      }
    }
    const staticIds = new Set(this.provider === undefined ? [] : this.provider.getModels().map((m) => m.id));
    let list = this.models.getModels(this.providerId);
    try {
      const available = await this.models.getAvailable(this.providerId);
      if (available.length) list = available;
    } catch {
      // 未配置凭据时 getAvailable 为空,退回完整目录
    }
    const source = (m: Model<Api>): PiModelInfo["source"] =>
      this.providerId === "openai-compatible" ? "custom" : refreshed && !staticIds.has(m.id) ? "dynamic" : "static";
    return { items: list.map((m) => modelInfo(m, source(m))), refreshed, refreshError };
  }

  /** 查模型;目录里没有(目录滞后或手填 id)时以同 provider 第一个模型为模板合成 */
  private pickModel(modelId: string | null | undefined): { model: Model<Api>; inCatalog: boolean } {
    const all = this.models.getModels(this.providerId);
    const id = modelId ?? this.spec.model ?? all[0]?.id;
    if (!id) throw new UpstreamError("model", "该模型服务没有可用模型,请填写模型 id");
    const found = this.models.getModel(this.providerId, id);
    if (found) return { model: found, inCatalog: true };
    const template = all[0];
    if (!template) throw new UpstreamError("model", `模型 ${id} 不在目录中,且该 provider 没有可参照的模型`);
    return { model: { ...template, id, name: id }, inCatalog: false };
  }

  /**
   * 测试:auth = 解析凭据 + 校验模型(不产生费用);completion = 再发一次最小补全(1 个输出 token)。
   * 失败抛 UpstreamError(kind 分类,消息已脱敏)。
   */
  async test(
    opts: { mode?: "auth" | "completion"; model?: string; signal?: AbortSignal } = {},
  ): Promise<{ source: string | null; model: string; inCatalog: boolean }> {
    const { source } = await this.resolveAuth();
    if (typeof this.provider.refreshModels === "function")
      await this.models
        .refresh({ providers: [this.providerId], allowNetwork: true, ...(opts.signal ? { signal: opts.signal } : {}) })
        .catch(() => undefined);
    const { model, inCatalog } = this.pickModel(opts.model);
    if (opts.mode === "completion") {
      let res: Awaited<ReturnType<MutableModels["completeSimple"]>>;
      try {
        res = await this.models.completeSimple(
          { ...model, maxTokens: Math.min(model.maxTokens || 16, 16) },
          { messages: [{ role: "user", content: "Reply with: ok", timestamp: Date.now() }] },
          {
            maxTokens: 16,
            ...(opts.signal ? { signal: opts.signal } : {}),
            timeoutMs: 30_000,
            maxRetries: 0,
            // OpenCode(Zen / Go)按会话路由,缺 x-opencode-session 会被拒;其余 provider 用它做缓存键,无副作用
            sessionId: `umt-upstream-test-${crypto.randomUUID()}`,
          },
        );
      } catch (err) {
        throw classifyUpstreamError(err);
      }
      if (res.stopReason === "error" || res.stopReason === "aborted")
        throw classifyUpstreamError(new Error(res.errorMessage ?? "模型服务返回错误"));
    }
    return { source, model: model.id, inCatalog };
  }

  /** 运行时装配:按设置裁剪上下文窗口 / 输出上限,思考档位夹到模型支持的范围 */
  async resolveRunModel(settings: ResolvedAgentSettings): Promise<ResolvedPiModel> {
    if (typeof this.provider.refreshModels === "function")
      await this.models.refresh({ providers: [this.providerId], allowNetwork: true }).catch(() => undefined);
    // 上游的 default_model 为准;未配置时用目录第一个模型(settings.model 只是记录用占位)
    const { model: m } = this.pickModel(this.spec.model ?? null);
    const contextWindow = Math.min(m.contextWindow || settings.contextBudgetTokens, settings.contextBudgetTokens);
    const model: Model<Api> = {
      ...m,
      contextWindow,
      maxTokens: Math.max(
        256,
        Math.min(m.maxTokens || settings.maxOutputTokens, settings.maxOutputTokens, Math.floor(contextWindow / 2)),
      ),
    };
    let thinkingLevel: AgentThinkingLevel = settings.thinkingLevel;
    try {
      thinkingLevel = clampThinkingLevel(model, settings.thinkingLevel) as AgentThinkingLevel;
    } catch {
      // 老模型无 thinking 元数据时保持设置值
    }
    return { models: this.models, model, thinkingLevel };
  }
}