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"
)

// 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
}

// 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 web-search endpoint (a SearXNG-compatible URL).
// WithSearchEndpoint("") clears it so web-search reports a clear error.
//
// 语义与既往逐字不变。**注意**：置空只是「端点为空」，工具仍在注册面里；
// 要「模型根本看不见它」请用 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）不是同一件事。
func WithSearchTimeout(d time.Duration) Option {
	return func(t *tools) {
		if d > 0 {
			t.searchTimeout = d
		}
	}
}

// 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(), nil
}

// 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 {
	file := newFileTool(t)
	set := []sdk.Tool{
		file,
		newCropTool(t, file),
		newHTTPTool(t),
	}
	// web-search 只有在**确实有可用端点**时才进注册面：端点显式关闭或非法时由
	// 组合根传 WithoutSearch，模型便根本看不见它。顺序与既往一致（http 之后、
	// web-fetch 之前），启用态下工具表的形状逐字不变。
	if !t.searchDisabled {
		set = append(set, newWebSearchTool(t))
	}
	return append(set,
		newWebFetchTool(t),
		newKnowledgeTool(t),
		newDocumentReadTool(t),
	)
}

// 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
}