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
}

// 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.
func WithSearchEndpoint(u string) Option {
	return func(t *tools) { t.searchURL = u }
}

// 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.
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)
	return []sdk.Tool{
		file,
		newCropTool(t, file),
		newHTTPTool(t),
		newWebSearchTool(t),
		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: 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(), t.allowedHosts); err != nil {
				return err
			}
			return nil
		},
	}
}