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 { const a = await this.ctx.settings.get("agent").catch(() => null); return { ...DEFAULT_WEB_SEARCH, ...(a?.webSearch ?? {}) }; } /** 设置里开启且存了 Key 才向模型提供 web.* */ async enabled(): Promise { 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 { 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( 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 { 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 { 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 }, }; } } }