UHD_Media_Tool/apps/server/src/modules/agent/capabilities-metadata.ts

import { randomUUID } from "node:crypto";
import { fetchPage, type HostResolver } from "@umt/adapters";
import {
  CapabilityError,
  type CapabilityRegistry,
  type HandlerContext,
  PiUpstreamClient,
  resolveAgentSettings,
} from "@umt/agent";
import {
  capabilitySpec,
  DEFAULT_WEB_SEARCH,
  type EnrichmentTrigger,
  type WebFetchOutput,
  type WebSearchOutput,
} from "@umt/contracts";
import { pgCacheFor } from "@umt/db";
import { resolveRunUpstream } from "../pi/index.ts";
import {
  AUTO_ENRICH_ACTOR,
  type EnrichDeps,
  EnrichError,
  type ModelCompletion,
  suggestEnrichment,
} from "../sources-identity/service/enrich.ts";
import { robotsAllows } from "../sources-identity/service/enrich-sources.ts";
import type { ServerContext } from "../types.ts";
import { agentSettings } from "./runner.ts";

// META-1 metadata.enrich(propose,low,replay safe):TMDB 查不到的作品按标题查网页出带证据的字段建议,不写作品表。
// 网络走 web.search / web.fetch 能力(同一注册中心:白名单 = 本次补全自己的搜索结果,审计照常记账);
// 模型用默认 Agent 上游做一次补全(fake 上游或没有上游时按规则抽取)。应用建议在控制台 /identity-pool/{id}/enrich/apply。

export interface MetadataCapabilityOptions {
  /** 模型调用(测试注入 fake);返回 null = 按规则抽取 */
  complete?: EnrichDeps["complete"];
  /** 读 robots.txt 用的 fetch / DNS 解析(测试注入 mock) */
  fetch?: typeof fetch;
  resolveHost?: HostResolver;
}

const ROBOTS_REGION = "enrich.robots";
const ROBOTS_TTL_SEC = 86400;

/**
 * META-2 robots.txt:按「协议 + 主机」缓存 1 天(kv_cache region enrich.robots)。
 * 404 / 其它 4xx = 没有限制;5xx、超时、网络错误 = 当作不允许(这次放弃该来源,不重试)。
 */
export async function robotsAllowedFor(
  ctx: Pick<ServerContext, "db">,
  url: string,
  opts: { fetch?: typeof fetch; resolveHost?: HostResolver; signal?: AbortSignal } = {},
): Promise<boolean> {
  let origin: string;
  try {
    origin = new URL(url).origin;
  } catch {
    return false;
  }
  const cache = pgCacheFor(ctx.db);
  let entry = await cache.get<{ txt: string | null; ok: boolean }>(ROBOTS_REGION, origin);
  if (!entry) {
    try {
      const p = await fetchPage(`${origin}/robots.txt`, {
        maxChars: 100_000,
        timeoutMs: 8000,
        ...(opts.fetch ? { fetch: opts.fetch } : {}),
        ...(opts.resolveHost ? { resolve: opts.resolveHost } : {}),
        ...(opts.signal ? { signal: opts.signal } : {}),
      });
      entry = { txt: p.text, ok: true };
    } catch (e) {
      const status = (e as { detail?: { status?: number } }).detail?.status;
      entry = status && status >= 400 && status < 500 ? { txt: null, ok: true } : { txt: null, ok: false };
    }
    await cache.set(ROBOTS_REGION, origin, entry, { ttl: ROBOTS_TTL_SEC, maxsize: 2000 });
  }
  if (!entry.ok) return false;
  return entry.txt ? robotsAllows(entry.txt, url) : true;
}

/** 默认 Agent 上游做一次非流式补全;fake 上游 / 未配置上游返回 null */
export async function upstreamComplete(
  ctx: ServerContext,
  system: string,
  user: string,
  signal?: AbortSignal,
): Promise<ModelCompletion | null> {
  const up = await resolveRunUpstream(ctx).catch(() => null);
  if (!up || up.row.provider === "fake") return null;
  const s = await ctx.settings.all();
  const settings = resolveAgentSettings(agentSettings(s, {}, false, up));
  const client = new PiUpstreamClient(up.spec, up.credentials);
  const { models, model } = await client.resolveRunModel(settings);
  const maxTokens = Math.max(512, Math.min(model.maxTokens || 2048, 2048));
  const res = await models.completeSimple(
    { ...model, maxTokens },
    { systemPrompt: system, messages: [{ role: "user", content: user, timestamp: Date.now() }] },
    {
      maxTokens,
      ...(signal ? { signal } : {}),
      timeoutMs: 90_000,
      maxRetries: 1,
      // OpenCode(Zen / Go)按会话路由,缺 x-opencode-session 会被拒
      sessionId: `umt-enrich-${randomUUID()}`,
    },
  );
  if (res.stopReason === "error" || res.stopReason === "aborted")
    throw new Error(res.errorMessage ?? "模型服务返回错误");
  const text = res.content
    .filter((b): b is { type: "text"; text: string } => b.type === "text")
    .map((b) => b.text)
    .join("");
  return {
    text,
    model: res.responseModel ?? res.model ?? null,
    inputTokens: res.usage?.input ?? 0,
    outputTokens: res.usage?.output ?? 0,
  };
}

function triggerOf(c: HandlerContext): EnrichmentTrigger {
  if (c.actorId === AUTO_ENRICH_ACTOR) return "auto";
  return c.source === "agent" || c.agentRunId ? "agent" : "manual";
}

export function registerMetadataCapabilities(
  registry: CapabilityRegistry,
  ctx: ServerContext,
  opts: MetadataCapabilityOptions = {},
): void {
  registry.register(capabilitySpec("metadata.enrich"), async (input, c) => {
    // 网页搜索没开就不占预算
    const agent = await ctx.settings.get("agent").catch(() => null);
    const ws = { ...DEFAULT_WEB_SEARCH, ...(agent?.webSearch ?? {}) };
    if (ws.provider === "none" || !ws.apiKeySecretRef || !registry.has("web.search"))
      throw new CapabilityError("precondition", "网页搜索未开启，在设置 → 智能体 → 网页搜索里开启后再补全");
    // 内部调用 web.* 用本次补全自己的调用方身份:白名单与 Agent 会话的互不覆盖
    const sub = (scope: string) => ({
      actorId: scope,
      taskId: c.taskId,
      scopeIds: [],
      signal: c.signal,
      source: "task" as const,
    });
    const deps: EnrichDeps = {
      async search(query, scope) {
        const out = (await registry.invoke("web.search", { query }, sub(scope))) as WebSearchOutput;
        return out.results.map((r) => ({ title: r.title || null, url: r.url, snippet: r.snippet || null }));
      },
      async fetch(url, scope) {
        const p = (await registry.invoke("web.fetch", { url }, sub(scope))) as WebFetchOutput;
        return { url: p.url, finalUrl: p.finalUrl, title: p.title, text: p.text, image: p.image ?? null };
      },
      complete: opts.complete ?? ((system, user, signal) => upstreamComplete(ctx, system, user, signal)),
      robotsAllowed: (url, signal) =>
        robotsAllowedFor(ctx, url, {
          ...(opts.fetch ? { fetch: opts.fetch } : {}),
          ...(opts.resolveHost ? { resolveHost: opts.resolveHost } : {}),
          ...(signal ? { signal } : {}),
        }),
    };
    try {
      return await suggestEnrichment(ctx, deps, input, {
        trigger: triggerOf(c),
        createdBy: c.actorId,
        signal: c.signal,
      });
    } catch (e) {
      if (e instanceof EnrichError) throw new CapabilityError(e.code, e.message);
      throw e;
    }
  });
}