Cyrene_v2/plugins/tools/io/tools.go

package io

import (
	"fmt"
	"log/slog"
	"net/http"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/platform"
	"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
	"git.yeij.top/AskaEth/Cyrene_v2/plugins/tools/io/search"
)

// source tags every log record emitted by this package (project rule §10).
const source = "plugins/tools/io"

// Default size, timeout and result limits. Named constants only — no bare
// magic numbers (project rule §10).
const (
	defaultMaxReadBytes   = 5 << 20 // 5 MiB file read cap
	defaultMaxWriteBytes  = 5 << 20 // 5 MiB file write cap
	defaultMaxHTTPBytes   = 1 << 20 // 1 MiB http response cap
	defaultMaxFetchBytes  = 2 << 20 // 2 MiB web-fetch cap
	defaultHTTPTimeout    = 15 * time.Second
	defaultMaxResults     = 10
	defaultMaxListEntries = 200
	maxRedirects          = 3
	filePerm              = 0o644
	dirPerm               = 0o755
)

// tools is the constructed, shared state for the whole tool set. A single HTTP
// client is shared by http/web-search/web-fetch so the connection pool is
// unified (project rule §10); it is never a package-level global.
type tools struct {
	cfg           *config.Config
	client        *http.Client
	allowedHosts  []string
	searchURL     string
	retriever     KnowledgeRetriever
	fileWhitelist []string
	maxReadBytes  int64
	maxWriteBytes int64
	maxHTTPBytes  int64
	maxFetchBytes int64
	httpTimeout   time.Duration
	maxResults    int
	maxList       int
	log           *slog.Logger

	// searchDisabled 让 web-search **根本不进注册面**（切片 WS1-A 的诚实性要求）：
	// 拿不到可用端点时，不该给模型一把拎起来必报错的锤子。见 WithoutSearch。
	searchDisabled bool
	// searchTimeout 是 web-search 单次请求的超时覆盖（0 = 沿用本套工具的
	// httpTimeout，也就是 io 包的既有超时 —— 这里刻意不另造默认值）。
	searchTimeout time.Duration
	// searchProviders 是组合根给的各家接线（WS1-B：provider 由配置选定）；
	// searchPrimary / searchFallback 是主 provider 与降级链。
	searchProviders []SearchProviderSpec
	searchPrimary   string
	searchFallback  []string
	// searchSelected 报告组合根是否**显式**选过 provider：只有显式选择才做启动期
	// fail-fast 校验（io 包的默认构造路径保留 WS1-A 的历史行为：工具在、调用时报错）。
	searchSelected bool
	// searchMaxResults 是 web-search 自己的条数上限（0 = 沿用 maxResults）。
	// 与 WithMaxResults 分开，是因为后者会一起改到 knowledge 工具的条数。
	searchMaxResults int
}

// SearchProviderSpec 是组合根给 web-search 的**一家** provider 的接线（一家一行）。
//
// Name 必须是 plugins/tools/io/search 的接入位（search.Known）；Configured 由组合根
// 用 search.Configured 判定后显式给出 —— 未配置的家不进注册表、不出现在可用清单里
// （design/27 §1.2 的 fail-closed）。
type SearchProviderSpec struct {
	Name       string
	APIKey     string
	BaseURL    string
	Configured bool
}

// Option customises the tool set built by New.
type Option func(*tools)

// WithHTTPTimeout sets the per-request timeout. Non-positive values are ignored.
func WithHTTPTimeout(d time.Duration) Option {
	return func(t *tools) {
		if d > 0 {
			t.httpTimeout = d
		}
	}
}

// WithHTTPClient injects the client shared by the network tools. When nil (the
// default) New builds one with a unified timeout and redirect SSRF guard.
func WithHTTPClient(c *http.Client) Option {
	return func(t *tools) {
		if c != nil {
			t.client = c
		}
	}
}

// WithAllowedHosts adds hosts to the SSRF allowlist. Hosts here bypass the
// loopback/private deny rules (they are the "host whitelist" side of the SSRF
// policy), letting the agent reach an internal search endpoint or the runtime.
func WithAllowedHosts(hosts ...string) Option {
	return func(t *tools) { t.allowedHosts = append(t.allowedHosts, hosts...) }
}

// WithSearchEndpoint sets the **searxng provider's** base endpoint (a SearXNG-compatible
// URL). WithSearchEndpoint("") clears it so web-search reports a clear error.
//
// 语义与既往逐字不变（design/27 §4：现有 WithSearchEndpoint 保留为 searxng provider 的
// 配置）。组合根走配置路径时用 WithSearchProviders 给出 searxng 的端点与可用性（单一
// 事实来源）；本选项保留给 io 包的既有调用方与测试。
//
// **注意**：置空只是「端点为空」，工具仍在注册面里；
// 要「模型根本看不见它」请用 WithoutSearch（组合根在端点显式关闭 / 非法时的用法）。
func WithSearchEndpoint(u string) Option {
	return func(t *tools) { t.searchURL = u }
}

// WithoutSearch removes the web-search tool from the built set entirely: it is
// not registered, so the model neither sees it nor wastes a round on it. This is
// the honest state for "no usable search endpoint" (WS1-A): a registered tool
// that can only ever error is worse than an absent one.
//
// 它只影响 web-search，其余工具照常构建。
func WithoutSearch() Option {
	return func(t *tools) { t.searchDisabled = true }
}

// WithSearchTimeout overrides the timeout of a single web-search call. Zero or
// negative means "use the set's existing HTTP timeout" — the io package does not
// define a second default for search (project rule §10: one default per knob).
//
// 它只约束 web-search 这一次请求：与 WithHTTPTimeout（整套工具共用的客户端
// 超时，会同时改到 http / web-fetch）不是同一件事。WS1-B 起它按**每一次上游尝试**
// 生效（降级链里每家都各有自己的预算）。
func WithSearchTimeout(d time.Duration) Option {
	return func(t *tools) {
		if d > 0 {
			t.searchTimeout = d
		}
	}
}

// WithSearchProviders 声明 web-search 可用的各家接线（WS1-B，design/27 §3）。
// 未传时退回 WS1-A 的既有路径：只有 searxng 一家（端点来自 WithSearchEndpoint）。
//
// 传进来的顺序即「可用清单」的展示顺序；Configured=false 的家不进注册表。
func WithSearchProviders(specs ...SearchProviderSpec) Option {
	return func(t *tools) { t.searchProviders = append(t.searchProviders, specs...) }
}

// WithSearchSelection 选定主 provider 与降级链（primary + fallback，按序）。
// 它同时打开**启动期 fail-fast 校验**：主 provider 未登记（如 tavily）或未配置（缺 key）
// 时 New 返回错误，绝不静默回落到兜底（design/27 §3）。
func WithSearchSelection(primary string, fallback []string) Option {
	return func(t *tools) {
		t.searchPrimary = primary
		t.searchFallback = append([]string(nil), fallback...)
		t.searchSelected = true
	}
}

// WithSearchMaxResults overrides web-search's own result cap (WEB_SEARCH_MAX_RESULTS).
// Non-positive values are ignored (the set's shared cap applies).
func WithSearchMaxResults(n int) Option {
	return func(t *tools) {
		if n > 0 {
			t.searchMaxResults = n
		}
	}
}

// WithKnowledgeRetriever wires the knowledge tool to a runtime memory retriever.
// WithKnowledgeRetriever(nil) makes knowledge report "no knowledge base".
func WithKnowledgeRetriever(r KnowledgeRetriever) Option {
	return func(t *tools) { t.retriever = r }
}

// WithFilePathWhitelist adds extra directories the file tool may access, in
// addition to the configured DataDir.
func WithFilePathWhitelist(paths ...string) Option {
	return func(t *tools) { t.fileWhitelist = append(t.fileWhitelist, paths...) }
}

// WithReadLimit overrides the file read cap (bytes). Non-positive values are
// ignored.
func WithReadLimit(n int64) Option {
	return func(t *tools) {
		if n > 0 {
			t.maxReadBytes = n
		}
	}
}

// WithWriteLimit overrides the file write cap (bytes). Non-positive values are
// ignored.
func WithWriteLimit(n int64) Option {
	return func(t *tools) {
		if n > 0 {
			t.maxWriteBytes = n
		}
	}
}

// WithHTTPResponseLimit overrides the http response cap (bytes). Non-positive
// values are ignored.
func WithHTTPResponseLimit(n int64) Option {
	return func(t *tools) {
		if n > 0 {
			t.maxHTTPBytes = n
		}
	}
}

// WithMaxResults caps the number of results returned by web-search and
// knowledge. Non-positive values are ignored.
func WithMaxResults(n int) Option {
	return func(t *tools) {
		if n > 0 {
			t.maxResults = n
		}
	}
}

// WithLogger injects the base logger. When nil the package builds one tagged
// with source "plugins/tools/io".
func WithLogger(l *slog.Logger) Option {
	return func(t *tools) {
		if l != nil {
			t.log = l
		}
	}
}

// New builds the full configured tool set over cfg and opts. It returns an error
// only when cfg is nil. A nil retriever or empty search endpoint is allowed: the
// corresponding tool reports a clear runtime error rather than failing startup.
//
// 组合根在「拿不到可用搜索端点」时应传 WithoutSearch，让 web-search **不进注册面**
// （模型看不见它），而不是塞一个空端点进去等着每次调用报错（切片 WS1-A）。
func New(cfg *config.Config, opts ...Option) ([]sdk.Tool, error) {
	if cfg == nil {
		return nil, fmt.Errorf("io: new: nil config")
	}
	t := &tools{
		cfg:           cfg,
		maxReadBytes:  defaultMaxReadBytes,
		maxWriteBytes: defaultMaxWriteBytes,
		maxHTTPBytes:  defaultMaxHTTPBytes,
		maxFetchBytes: defaultMaxFetchBytes,
		httpTimeout:   defaultHTTPTimeout,
		maxResults:    defaultMaxResults,
		maxList:       defaultMaxListEntries,
		log:           logger.WithSource(logger.New(logger.Options{}), source),
	}
	for _, opt := range opts {
		if opt != nil {
			opt(t)
		}
	}
	if t.client == nil {
		t.client = newHTTPClient(t)
	}
	return t.build()
}

// Tools returns the tool set built from platform defaults. Production wiring
// should prefer New(cfg, opts...) so the tools honor the real configuration.
func Tools() []sdk.Tool {
	set, err := New(&config.Config{DataDir: platform.DefaultDataDir()})
	if err != nil {
		return nil
	}
	return set
}

// Register registers the default tool set (from Tools) with reg. It returns
// sdk.ErrNilRegistrar when reg is nil and stops at the first registration error.
func Register(reg sdk.Registrar) error {
	if reg == nil {
		return sdk.ErrNilRegistrar
	}
	for _, t := range Tools() {
		if err := sdk.RegisterTool(reg, t); err != nil {
			return fmt.Errorf("io: register %q: %w", t.Name(), err)
		}
	}
	return nil
}

// build assembles the tools sharing this set's client, limits and logger.
func (t *tools) build() ([]sdk.Tool, error) {
	file := newFileTool(t)
	set := []sdk.Tool{
		file,
		newCropTool(t, file),
		newHTTPTool(t),
	}
	// web-search 只有在**确实有可用后端**时才进注册面：端点显式关闭或选定 provider
	// 未登记/未配置时由组合根传 WithoutSearch，模型便根本看不见它。顺序与既往一致
	// （http 之后、web-fetch 之前），启用态下工具表的形状逐字不变。
	if !t.searchDisabled {
		ws, err := newWebSearchTool(t)
		if err != nil {
			// 启动期 fail-fast（design/27 §3）：选定的 provider 有问题就不许静默回落。
			return nil, err
		}
		set = append(set, ws)
	}
	return append(set,
		newWebFetchTool(t),
		newKnowledgeTool(t),
		newDocumentReadTool(t),
	), nil
}

// searchWires 汇总要交给搜索注册表的接线（WS1-B）：
//
//   - 组合根给了接线表（WithSearchProviders）⇒ 用它，是否可用由 Configured 决定；
//   - 没给 ⇒ 退回 WS1-A 的既有路径：只有 searxng 一家，且**始终登记**（哪怕端点是
//     空串）—— 那是历史行为：工具在注册面里、调用时如实报「no search endpoint」。
func (t *tools) searchWires() []search.Wire {
	if len(t.searchProviders) == 0 {
		return []search.Wire{{Name: search.NameSearXNG, BaseURL: t.searchURL, Configured: true}}
	}
	wires := make([]search.Wire, 0, len(t.searchProviders))
	for _, s := range t.searchProviders {
		wires = append(wires, search.Wire{
			Name:       s.Name,
			APIKey:     s.APIKey,
			BaseURL:    s.BaseURL,
			Configured: s.Configured,
		})
	}
	return wires
}

// newHTTPClient builds a client with a unified timeout and an SSRF-guarded
// redirect policy that refuses to follow into private/loopback targets.
func newHTTPClient(t *tools) *http.Client {
	return &http.Client{
		Timeout:       t.httpTimeout,
		CheckRedirect: redirectGuard(t.allowedHosts),
	}
}

// redirectGuard 造 SSRF 重定向守卫：最多跳 maxRedirects 次，且每次跳转的目标都要
// 过同一份判定（跳去内网、回环一样要拦）。判定所用的白名单是参数给的这一份，
// 因此调用方能让它按**自己**加宽后的名单判定，而不牵动共享白名单。
func redirectGuard(allowedHosts []string) func(*http.Request, []*http.Request) error {
	return func(req *http.Request, via []*http.Request) error {
		if len(via) >= maxRedirects {
			return fmt.Errorf("io: too many redirects")
		}
		if req.URL == nil {
			return ErrSSRF
		}
		if _, err := validateRequestURL(req.Context(), req.URL.String(), allowedHosts); err != nil {
			return err
		}
		return nil
	}
}

// clientFor 返回「白名单为 allowedHosts」时某个工具该用的客户端。
//
// allowedHosts 与共享白名单**长度相同**即代表没有加宽（withSearchEndpointHost 只会
// 追加、绝不改写已有条目），此时直接复用共享客户端；确有加宽时返回一份**副本**，
// 把重定向守卫换成按加宽名单判定的那一份 —— 副本共享同一个 Transport，连接池不变。
//
// 这条分工的意义：让 web-search 认自己配置里的端点主机（含本机兜底的回环地址），
// 同时**不**污染 http / web-fetch 共用的那份白名单与那个客户端。
func (t *tools) clientFor(allowedHosts []string) *http.Client {
	if len(allowedHosts) == len(t.allowedHosts) {
		return t.client
	}
	client := *t.client
	client.CheckRedirect = redirectGuard(allowedHosts)
	return &client
}