UHD_Media_Tool/apps/server/src/modules/websearch/service.ts

import { createHash } from "node:crypto";
import {
  createBraveClient,
  fetchPage,
  type HostResolver,
  WebSearchError,
  type WebSearchErrorKind,
} from "@umt/adapters";
import {
  DEFAULT_WEB_SEARCH,
  UNTRUSTED_WEB_NOTICE,
  type WebFetchOutput,
  type WebSearchFreshness,
  type WebSearchOutput,
  type WebSearchResult,
  type WebSearchSettings,
  type WebSearchTestResult,
} from "@umt/contracts";
import { cachedCall, pgCacheFor, pgTokenBucketFor } from "@umt/db";
import { externalFetch } from "../../egress.ts";
import { storageSettings } from "../../shared-cache.ts";
import type { ServerContext } from "../types.ts";

/**
 * WEB-1 网页搜索服务:Brave 客户端之上的 Key 读取、每分钟上限(跨进程令牌桶)、PG 缓存(kv_cache region websearch,TTL = storage.cacheTtl.webSearchDays)、
 * 以及 web.fetch 的白名单——按调用方(Agent 会话 / 其它调用者)记录最近一次 web.search 返回的 URL,只允许读取这些 URL。
 */

export interface WebSearchModuleOptions {
  /** 搜索与网页读取用的 fetch(测试注入 mock) */
  fetch?: typeof fetch;
  /** Brave 基址(测试指向 mock) */
  braveBaseUrl?: string;
  /** DNS 解析(测试注入);缺省系统解析 */
  resolveHost?: HostResolver;
}

export const WEBSEARCH_CACHE_REGION = "websearch";
/** 白名单:key = 调用方,value = 最近一次搜索的 URL */
export const WEBSEARCH_ALLOW_REGION = "websearch.allow";
const SEARCH_CACHE_MAX = 1000;
/** 白名单保留 6 小时(会话时长上限远小于此) */
const ALLOW_TTL_SEC = 6 * 3600;
/** 取不到令牌时最多等多久 */
const RATE_WAIT_MS = 3000;
export const TEST_QUERY = "MoviePilot";

export class WebSearchServiceError extends Error {
  constructor(
    readonly kind: WebSearchErrorKind | "disabled" | "not_allowed",
    message: string,
  ) {
    super(message);
    this.name = "WebSearchServiceError";
  }
}

/** URL 比较用的规范形式(去掉 #片段) */
export function canonicalUrl(raw: string): string | null {
  try {
    const u = new URL(raw);
    u.hash = "";
    return u.href;
  } catch {
    return null;
  }
}

/** web.fetch 白名单的归属:Agent 会话按 run,其它调用方按 actor */
export const allowScopeKey = (c: { agentRunId?: string | undefined; actorId: string }) =>
  c.agentRunId ? `run:${c.agentRunId}` : `actor:${c.actorId}`;

export class WebSearchService {
  constructor(
    private readonly ctx: ServerContext,
    private readonly opts: WebSearchModuleOptions = {},
  ) {}

  async settings(): Promise<WebSearchSettings> {
    const a = await this.ctx.settings.get("agent").catch(() => null);
    return { ...DEFAULT_WEB_SEARCH, ...(a?.webSearch ?? {}) };
  }

  /** 设置里开启且存了 Key 才向模型提供 web.* */
  async enabled(): Promise<boolean> {
    const s = await this.settings();
    return s.provider !== "none" && !!s.apiKeySecretRef;
  }

  private async client(s: WebSearchSettings) {
    if (s.provider === "none")
      throw new WebSearchServiceError("disabled", "网页搜索未开启，在设置 → 智能体 → 网页搜索里开启");
    if (!s.apiKeySecretRef) throw new WebSearchServiceError("auth", "未保存搜索服务 API Key");
    const apiKey = await this.ctx.secrets.resolve(s.apiKeySecretRef);
    return createBraveClient({
      apiKey,
      ...(this.opts.braveBaseUrl ? { baseUrl: this.opts.braveBaseUrl } : {}),
      // 外网请求:显式按环境变量走代理(CF-2 出口表)
      fetch: this.opts.fetch ?? externalFetch,
    });
  }

  private async takeToken(s: WebSearchSettings, signal?: AbortSignal) {
    const rpm = s.requestsPerMinute;
    const r = await pgTokenBucketFor(this.ctx.db).acquire(`websearch:${s.provider}`, 1, {
      capacity: rpm,
      refillPerSec: rpm / 60,
      waitMs: RATE_WAIT_MS,
      ...(signal ? { signal } : {}),
    });
    if (!r.ok) throw new WebSearchServiceError("rate", `网页搜索已达每分钟上限（${rpm} 次），稍后再试`);
  }

  /**
   * 搜索:命中缓存直接返回(不计上限);否则扣令牌后查询并缓存 10 分钟。
   * scopeKey 给了就把结果 URL 记为该调用方的 web.fetch 白名单(替换上一次的)。
   */
  async search(
    input: { query: string; count?: number | undefined; freshness?: WebSearchFreshness | undefined },
    o: { scopeKey?: string; fresh?: boolean; signal?: AbortSignal } = {},
  ): Promise<WebSearchOutput & { durationMs: number }> {
    const s = await this.settings();
    const client = await this.client(s);
    const count = Math.min(input.count ?? s.maxResults, s.maxResults);
    const params = {
      q: input.query.trim(),
      count,
      country: s.country === "ALL" ? null : s.country,
      search_lang: s.searchLang,
      safesearch: s.safesearch,
      ...(input.freshness ? { freshness: input.freshness } : {}),
    };
    const key = createHash("sha256")
      .update(
        JSON.stringify([s.provider, params.q, count, s.country, s.searchLang, s.safesearch, input.freshness ?? null]),
      )
      .digest("hex")
      .slice(0, 40);
    let hit = true;
    const started = Date.now();
    const results = await cachedCall<WebSearchResult[]>(
      pgCacheFor(this.ctx.db),
      key,
      {
        region: WEBSEARCH_CACHE_REGION,
        ttl: (await storageSettings(this.ctx)).cacheTtl.webSearchDays * 86_400,
        maxsize: SEARCH_CACHE_MAX,
        ...(o.fresh ? { fresh: true } : {}),
      },
      async () => {
        hit = false;
        await this.takeToken(s, o.signal);
        const r = await client.search({ ...params, ...(o.signal ? { signal: o.signal } : {}) });
        return r.results;
      },
    );
    if (o.scopeKey) {
      const urls = results.map((r) => canonicalUrl(r.url)).filter((u): u is string => !!u);
      await pgCacheFor(this.ctx.db).set(
        WEBSEARCH_ALLOW_REGION,
        o.scopeKey,
        { urls, query: params.q },
        { ttl: ALLOW_TTL_SEC, maxsize: 5000 },
      );
    }
    return {
      trust: "untrusted",
      notice: UNTRUSTED_WEB_NOTICE,
      query: params.q,
      results,
      cached: hit,
      durationMs: Date.now() - started,
    };
  }

  /** 读取网页:必须是该调用方最近一次搜索返回过的 URL;网络层再拒绝内网与非 http(s) */
  async fetch(url: string, o: { scopeKey: string; signal?: AbortSignal }): Promise<WebFetchOutput> {
    const s = await this.settings();
    if (s.provider === "none")
      throw new WebSearchServiceError("disabled", "网页搜索未开启，在设置 → 智能体 → 网页搜索里开启");
    const want = canonicalUrl(url.trim());
    const allow = await pgCacheFor(this.ctx.db).get<{ urls: string[] }>(WEBSEARCH_ALLOW_REGION, o.scopeKey);
    if (!want || !allow?.urls.includes(want))
      throw new WebSearchServiceError(
        "not_allowed",
        allow ? "只能读取最近一次 web.search 结果里的链接" : "先用 web.search 搜索，再读取结果里的链接",
      );
    const page = await fetchPage(want, {
      fetch: this.opts.fetch ?? externalFetch,
      ...(this.opts.resolveHost ? { resolve: this.opts.resolveHost } : {}),
      ...(o.signal ? { signal: o.signal } : {}),
    });
    return { trust: "untrusted", notice: UNTRUSTED_WEB_NOTICE, ...page };
  }

  /** 设置页「测试」:用已保存的 Key 实际查一次(不走缓存) */
  async test(query?: string): Promise<WebSearchTestResult> {
    const s = await this.settings();
    const q = query?.trim() || TEST_QUERY;
    const started = Date.now();
    try {
      const r = await this.search({ query: q, count: Math.min(3, s.maxResults) }, { fresh: true });
      return {
        ok: true,
        provider: s.provider,
        query: q,
        count: r.results.length,
        durationMs: Date.now() - started,
        results: r.results.map(({ title, url, source }) => ({ title, url, source })),
        error: null,
      };
    } catch (err) {
      const e = err as WebSearchError | WebSearchServiceError;
      if (!(e instanceof WebSearchError) && !(e instanceof WebSearchServiceError)) throw err;
      return {
        ok: false,
        provider: s.provider,
        query: q,
        count: 0,
        durationMs: Date.now() - started,
        results: [],
        error: { code: e.kind, message: e.message },
      };
    }
  }
}