Cyrene_v2/plugins/tools/io/web_search.go

package io

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"net/url"
	"strings"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
	"git.yeij.top/AskaEth/Cyrene_v2/plugins/tools/io/search"
)

// Compile-time assertion that webSearchTool implements sdk.Tool.
var _ sdk.Tool = (*webSearchTool)(nil)

// webSearchTool 是**一把**能力面对模型的工具（design/27 §1.1：能力 ≠ 单一后端，
// 换 provider 不改工具 schema）。它把「选谁当后端」交给注入的注册表，
// 并按 主 provider → 降级链 依次尝试（design/27 §3）。
type webSearchTool struct {
	client   *http.Client
	registry *search.Registry
	// primary 是主 provider 名；fallback 是降级链（按序）。
	primary  string
	fallback []string
	// allowedHosts 是**本工具自己**那份 SSRF 白名单（含配置里各 provider 的主机），
	// 不是 http / web-fetch 共用的那份（见 newWebSearchTool 的说明）。
	allowedHosts []string
	// endpoint 是 searxng provider 的端点（WS1-A 的 WEB_SEARCH_ENDPOINT 语义，
	// 保留为字段以便既有测试与诊断继续可用）。
	endpoint   string
	maxBytes   int64
	maxResults int
	log        *slog.Logger
	// timeout 是**每次**上游尝试的超时预算（0 = 沿用共享客户端自己的超时）。
	// 按「每次尝试」而不是「整个调用」计，降级链才真的有意义：主家超时后，
	// 备家仍有自己的预算可用。
	timeout time.Duration
}

// newWebSearchTool 构造工具：汇总接线 → 校验选定 → 建注册表 → 按各家主机加宽
// **本工具自己**的 SSRF 白名单。
//
// SSRF 说明（沿用 WS1-A 的口径）：端点是**运维配置**、不是模型输入（模型的 query 只
// 进查询串），所以把配置里的 provider 主机并入 web-search 自己的白名单**不会**给模型
// 打开新的目标面——工具能去的主机恒等于配置里的那几个。http / web-fetch 共用白名单
// 保持不变（TestWebSearch_DoesNotWidenSharedAllowlist 钉住）。
func newWebSearchTool(t *tools) (*webSearchTool, error) {
	wires := t.searchWires()
	primary := strings.TrimSpace(t.searchPrimary)
	if primary == "" {
		primary = search.NameSearXNG
	}
	// 校验只在「配置路径」生效（组合根显式给了选择或接线表）：io 包的既有默认构造
	// （WS1-A 之前就在的调用方与测试）刻意保留「工具在、调用时如实报错」的历史行为。
	if t.searchSelected || len(t.searchProviders) > 0 {
		if err := checkSearchSelection(primary, wires); err != nil {
			return nil, err
		}
	}
	allowedHosts := t.allowedHosts
	for _, w := range wires {
		allowedHosts = withSearchEndpointHost(allowedHosts, w.BaseURL)
	}
	client := t.clientFor(allowedHosts)
	maxResults := t.maxResults
	if t.searchMaxResults > 0 {
		maxResults = t.searchMaxResults
	}
	return &webSearchTool{
		client:       client,
		registry:     search.Build(wires, client, t.maxHTTPBytes, t.searchGuard(allowedHosts)),
		primary:      primary,
		fallback:     t.searchFallback,
		allowedHosts: allowedHosts,
		// endpoint 是**searxng provider 实际用的端点**（诊断/既有测试用）：从接线表里
		// 取，因此无论走 WithSearchEndpoint 还是组合根的接线表，它都反映真实取值。
		endpoint:   searxngEndpoint(wires),
		maxBytes:   t.maxHTTPBytes,
		maxResults: maxResults,
		log:        t.log,
		timeout:    t.searchTimeout,
	}, nil
}

// searxngEndpoint 从接线表里取出 searxng 那一条的基地址（没有则空串）。
func searxngEndpoint(wires []search.Wire) string {
	for _, w := range wires {
		if strings.EqualFold(w.Name, search.NameSearXNG) {
			return w.BaseURL
		}
	}
	return ""
}

// checkSearchSelection 是启动期校验（design/27 §3 的语义硬约束）：
//
//   - 名字**不认识**（写错了 / 本期没做，如 tavily）⇒ ErrUnknownProvider；
//   - 名字认识但**没配齐**（云家缺 key、searxng 缺端点）⇒ ErrNotConfigured。
//
// 两者都让启动显式失败，**绝不静默回落到兜底** —— 否则「我配了 Tavily」就是一句假话。
func checkSearchSelection(primary string, wires []search.Wire) error {
	if !search.Known(primary) {
		return fmt.Errorf("io: web-search: %w: %q (known: %s)",
			search.ErrUnknownProvider, primary, strings.Join(search.KnownNames(), ", "))
	}
	for _, w := range wires {
		if strings.EqualFold(w.Name, primary) && w.Configured {
			return nil
		}
	}
	return fmt.Errorf("io: web-search: %w: %q (its API key / base URL is missing)", search.ErrNotConfigured, primary)
}

// searchGuard 造**本工具自己**的 SSRF 守卫：复用 io 包既有的 validateRequestURL 与
// 「配置里各 provider 主机」那份加宽白名单（allowedHosts），因此
//   - 端点/基地址是**运维配置**（不是模型输入）⇒ 回环的本机兜底仍然可达；
//   - 模型只能改查询串，改不了主机 ⇒ 放行这几个主机不给模型开新的目标面；
//   - http / web-fetch 共用的那份白名单**不受影响**（见 clientFor 的分工）。
func (t *tools) searchGuard(allowedHosts []string) search.Guard {
	return func(ctx context.Context, rawURL string) error {
		if _, err := validateRequestURL(ctx, rawURL, allowedHosts); err != nil {
			// 守卫的判定已在 validateRequestURL 里，这里只补上是谁拦的（保持 %w 链）。
			return fmt.Errorf("io: web-search: %w", err)
		}
		return nil
	}
}

// Name implements sdk.Tool.
func (s *webSearchTool) Name() string { return "web-search" }

// Description implements sdk.Tool. 工具面对模型只有这一把，后端由运维配置决定。
func (s *webSearchTool) Description() string {
	return "Search the web through the configured search provider and return ranked results."
}

// Parameters implements sdk.Tool. Schema 与既往逐字一致（design/27 §4：不破坏契约）。
func (s *webSearchTool) Parameters() (map[string]any, error) {
	return map[string]any{
		"type": "object",
		"properties": map[string]any{
			"query": map[string]any{"type": "string", "description": "Search query."},
			"count": map[string]any{"type": "integer", "description": "Max results to return."},
		},
		"required": []any{"query"},
	}, nil
}

// Execute implements sdk.Tool.
func (s *webSearchTool) Execute(ctx context.Context, args any) (any, error) {
	m, ok := args.(map[string]any)
	if !ok {
		return nil, fmt.Errorf("io: web-search: args must be an object")
	}
	query := strings.TrimSpace(argString(m, "query"))
	if query == "" {
		return nil, fmt.Errorf("io: web-search: query is required")
	}
	n := s.maxResults
	if v, ok := argInt(m, "count"); ok && v > 0 {
		n = v
	}

	raw, provider, err := s.dispatch(ctx, query, n)
	if err != nil {
		return nil, err
	}
	// 归一化之后**统一做一次**后处理：域名过滤/去重/排序/截断（design/27 §4）。
	// ⚠️ 这里判的是「**归一化之后**确实一条都没有」，不是「上游没给链接」—— 后者曾被
	// Finalize 当成没有结果（2026-10-07 智谱真机缺陷），现在无链接的条目照原样保留。
	results := search.Finalize(raw, search.Options{Limit: n})
	if len(results) == 0 {
		// 空结果 ≠ 成功（design/27 §1.4）：如实报「一家都没查到」，绝不冒充成功。
		// 它同样**不**触发降级链前进（那是「不可用」的处置，见 dispatch）。
		return nil, fmt.Errorf("io: web-search: %w (provider=%s)", search.ErrNoResults, provider)
	}
	return map[string]any{
		"query":   query,
		"count":   len(results),
		"results": toWireResults(results),
		// truncated 沿用既有字段（成功路径恒为 false：超限在 provider 层就是错误）。
		"truncated": false,
		// provider 是**如实标注**：降级链里谁应答就写谁，绝不拿兜底冒充主力。
		"provider": provider,
	}, nil
}

// dispatch 按 主 provider → 降级链 依次尝试，返回结果与**实际应答的 provider 名**。
//
// 只有「该家不可用」才前进一格（unavailable：无 key / 鉴权失败 / 限流 / 额度 / 网络 /
// 上游错误）；「空结果」与「预算耗尽」**不**前进 —— 空是空，不是不可用（design/27 §1.4/§3）。
func (s *webSearchTool) dispatch(ctx context.Context, query string, count int) ([]search.Result, string, error) {
	names := make([]string, 0, len(s.fallback)+1)
	names = append(names, s.primary)
	names = append(names, s.fallback...)

	attempted := make([]string, 0, len(names))
	var lastErr error
	for _, name := range names {
		p, ok := s.registry.Lookup(name)
		if !ok {
			// 未配置/未登记的家在构造期就不在注册表里，链上直接跳过（不是错误）。
			s.log.Debug("web-search: provider not available, skipped",
				"provider", name, "available", s.registry.Names())
			continue
		}
		if err := ctx.Err(); err != nil {
			// 上游预算已耗尽，再开一格也没有意义。
			lastErr = err
			break
		}
		attempted = append(attempted, name)
		res, err := s.searchWith(p, ctx, query, count)
		if err == nil {
			return res, name, nil
		}
		lastErr = err
		if !searchUnavailable(err) {
			// 空结果 / 上下文取消：停在原地，如实返回。
			return nil, name, fmt.Errorf("io: web-search: %w", err)
		}
		// 错误文案**不含密钥**（密钥只在请求头里，从不进错误/日志）。
		s.log.Warn("web-search: provider unavailable, advancing fallback chain",
			"provider", name, "err", err)
	}
	if lastErr == nil {
		lastErr = search.ErrNotConfigured
	}
	if len(attempted) == 0 {
		return nil, "", fmt.Errorf("io: web-search: %w (no provider available)", lastErr)
	}
	return nil, "", fmt.Errorf("io: web-search: %w (tried: %s)", lastErr, strings.Join(attempted, ", "))
}

// searchWith 调一家 provider，并按**这一次尝试**的预算加超时。
//
// 「上游正常应答但一条都没查到」在这里就被点成 ErrNoResults（并**带上是谁答的**）——
// 它是「空」不是「成功」，也不是「该家不可用」：dispatch 据此**停在原地**、如实报出去，
// 而不是换一家假装就有结果了（design/27 §1.4/§3）。把这次判定放在这里，是为了让
// 「空结果不前进」这条规则只有一处，且真的被执行到。
func (s *webSearchTool) searchWith(p search.Provider, ctx context.Context, query string, count int) ([]search.Result, error) {
	callCtx := ctx
	if s.timeout > 0 {
		var cancel context.CancelFunc
		callCtx, cancel = context.WithTimeout(ctx, s.timeout)
		defer cancel()
	}
	res, err := p.Search(callCtx, search.Query{Text: query, Count: count})
	if err != nil {
		return nil, err
	}
	if len(res) == 0 {
		return nil, fmt.Errorf("%w (provider=%s)", search.ErrNoResults, p.Name())
	}
	return res, nil
}

// searchUnavailable 报告一个错误是否属于「该家不可用」。
//
// 上下文取消/超时**不算**：那是本次调用的预算花光了，换一家也一样没时间，
// 应该如实报出去而不是掩盖成「前进一格」。
func searchUnavailable(err error) bool {
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
		return false
	}
	return search.Unavailable(err)
}

// searchResult 是一条结果的**线格式**（既有字段一字不改；published/score 是本期追加的
// 可选字段 —— design/27 §4：只许追加，不许改名）。
type searchResult struct {
	Title string `json:"title"`
	// URL 是结果链接。⚠️ **可能是空串**：上游（如智谱 `search_std`）真机上经常不给 link，
	// 此时如实回空串（`title`/`content` 仍可用），**绝不**编一个 URL 出来。
	// 客户端按「空 url = 这条没有可点链接」处理，不要把它当成失败或省略字段。
	URL     string `json:"url"`
	Content string `json:"content"`
	Snippet string `json:"snippet"`
	// Published 是发布时间（各家口径不一：ISO 字符串或毫秒时间戳）。
	Published string `json:"published,omitempty"`
	// Score 是上游给的排序分（多数家不给，此时为 0 并被省略）。
	Score float64 `json:"score,omitempty"`
}

// toWireResults 把归一化结果映射成线格式（追加 provider/score 等字段，不改既有字段名）。
func toWireResults(results []search.Result) []searchResult {
	out := make([]searchResult, 0, len(results))
	for _, r := range results {
		out = append(out, searchResult{
			Title:     r.Title,
			URL:       r.Link,
			Content:   r.Content,
			Snippet:   r.Snippet,
			Published: r.Published,
			Score:     r.Score,
		})
	}
	return out
}

// withSearchEndpointHost 把**配置里给定的**provider 主机并入 web-search 自己的
// SSRF 白名单（返回新切片，不改动入参切片——组合根可能把它交给别的工具）。
// 端点解析不出主机、或主机已在名单里时原样返回。
//
// 理由见 newWebSearchTool：本机兜底端点必然是回环地址（127.0.0.1:38471），而
// SSRF 守卫默认拒绝回环；端点是运维配置而非模型输入，因此对它放行**不会**给模型
// 打开新的目标面——工具能去的主机恒等于配置里的那几个。
func withSearchEndpointHost(allowed []string, endpoint string) []string {
	host := searchEndpointHost(endpoint)
	if host == "" || isHostAllowed(host, allowed) {
		return allowed
	}
	out := make([]string, 0, len(allowed)+1)
	out = append(out, allowed...)
	return append(out, host)
}

// searchEndpointHost 取出 provider 基地址的主机名（去端口、去 IPv6 方括号）；解析不出来
// 或没有主机时返回空串。
func searchEndpointHost(endpoint string) string {
	u, err := url.Parse(strings.TrimSpace(endpoint))
	if err != nil {
		return ""
	}
	return hostOnly(u.Host)
}