feat(tools/io): 内置 IO/网络工具(file(路径安全)/http(SSRF防护)/web-search/web-fetch/knowledge)+ 大小截断
This commit is contained in:
1 parent
488cf9c79c
commit
3ce9489f8e
17 files changed
+2223
No files matched your search
@@ -0,0 +1,70 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
)
|
||||||
|
|
||||||
|
// argString returns the string value of args[key], or "" when absent or of
|
||||||
|
// another type.
|
||||||
|
func argString(args map[string]any, key string) string {
|
||||||
|
if args == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
s, _ := args[key].(string)
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
|
||||||
|
// argInt returns the integer value of args[key] and whether it was a number.
|
||||||
|
func argInt(args map[string]any, key string) (int, bool) {
|
||||||
|
if args == nil {
|
||||||
|
return 0, false
|
||||||
|
}
|
||||||
|
return toInt(args[key])
|
||||||
|
}
|
||||||
|
|
||||||
|
// argStringMap coerces a JSON object argument to a string map, accepting both
|
||||||
|
// map[string]any and map[string]string inputs, dropping non-string values.
|
||||||
|
func argStringMap(args map[string]any, key string) map[string]string {
|
||||||
|
out := map[string]string{}
|
||||||
|
if args == nil {
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
switch v := args[key].(type) {
|
||||||
|
case map[string]any:
|
||||||
|
for k, val := range v {
|
||||||
|
if s, ok := val.(string); ok {
|
||||||
|
out[k] = s
|
||||||
|
}
|
||||||
|
}
|
||||||
|
case map[string]string:
|
||||||
|
for k, val := range v {
|
||||||
|
out[k] = val
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// toInt converts a decoded-JSON number (float64 / json.Number) or a Go integer
|
||||||
|
// to an int, reporting whether the conversion succeeded.
|
||||||
|
func toInt(v any) (int, bool) {
|
||||||
|
switch x := v.(type) {
|
||||||
|
case int:
|
||||||
|
return x, true
|
||||||
|
case int32:
|
||||||
|
return int(x), true
|
||||||
|
case int64:
|
||||||
|
return int(x), true
|
||||||
|
case float32:
|
||||||
|
return int(x), true
|
||||||
|
case float64:
|
||||||
|
return int(x), true
|
||||||
|
case json.Number:
|
||||||
|
n, err := x.Int64()
|
||||||
|
if err != nil {
|
||||||
|
return 0, false
|
||||||
|
}
|
||||||
|
return int(n), true
|
||||||
|
default:
|
||||||
|
return 0, false
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,25 @@
|
|||||||
|
// Package io implements the built-in I/O and network tools for Cyrene v2
|
||||||
|
// (DevDocs/design/09-插件生态.md). It exposes file, http, web-search, web-fetch
|
||||||
|
// and knowledge as pkg/plugins/sdk.Tool values, wired to the process-internal
|
||||||
|
// config, logger and runtime/memory retriever.
|
||||||
|
//
|
||||||
|
// Build the full configured tool set with New(cfg, opts...). Tools() and
|
||||||
|
// Register(reg) are convenience wrappers that use platform defaults; production
|
||||||
|
// wiring should call New(cfg, opts...) so the tools honor the real DataDir,
|
||||||
|
// search endpoint and knowledge retriever.
|
||||||
|
//
|
||||||
|
// Security constraints
|
||||||
|
//
|
||||||
|
// - file: every path is cleaned and validated to stay inside the configured
|
||||||
|
// data directory or an explicit whitelist; ".." traversal and symlink escape
|
||||||
|
// are rejected; reads and writes are size-capped.
|
||||||
|
// - http: GET/POST only, timeouts, response truncation and SSRF protection
|
||||||
|
// (loopback/private/link-local addresses and localhost are refused unless
|
||||||
|
// the host is on an explicit allowlist).
|
||||||
|
// - web-search: requires a configured search endpoint, otherwise a clear error.
|
||||||
|
// - web-fetch: size- and content-type-limited readable text extraction.
|
||||||
|
// - knowledge: RAG lookup over the runtime memory retriever; a clear message
|
||||||
|
// when no knowledge base is configured.
|
||||||
|
//
|
||||||
|
// All tools are safe for concurrent use, as required by sdk.Tool.
|
||||||
|
package io
|
||||||
@@ -0,0 +1,311 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
stdio "io"
|
||||||
|
"io/fs"
|
||||||
|
"log/slog"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Compile-time assertion that fileTool implements sdk.Tool.
|
||||||
|
var _ sdk.Tool = (*fileTool)(nil)
|
||||||
|
|
||||||
|
// fileTool reads, writes, lists and searches files, confined to the configured
|
||||||
|
// data directory plus any explicit whitelist. Paths are cleaned and re-checked
|
||||||
|
// against those roots, so ".." traversal and symlink escapes are rejected.
|
||||||
|
type fileTool struct {
|
||||||
|
roots []string // absolute, cleaned allowed directories
|
||||||
|
maxRead int64
|
||||||
|
maxWrite int64
|
||||||
|
maxList int
|
||||||
|
log *slog.Logger
|
||||||
|
}
|
||||||
|
|
||||||
|
func newFileTool(t *tools) *fileTool {
|
||||||
|
roots := make([]string, 0, 1+len(t.fileWhitelist))
|
||||||
|
if t.cfg != nil && t.cfg.DataDir != "" {
|
||||||
|
roots = append(roots, absoluteClean(t.cfg.DataDir))
|
||||||
|
}
|
||||||
|
for _, w := range t.fileWhitelist {
|
||||||
|
roots = append(roots, absoluteClean(w))
|
||||||
|
}
|
||||||
|
return &fileTool{
|
||||||
|
roots: roots,
|
||||||
|
maxRead: t.maxReadBytes,
|
||||||
|
maxWrite: t.maxWriteBytes,
|
||||||
|
maxList: t.maxList,
|
||||||
|
log: t.log,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Name implements sdk.Tool.
|
||||||
|
func (f *fileTool) Name() string { return "file" }
|
||||||
|
|
||||||
|
// Description implements sdk.Tool.
|
||||||
|
func (f *fileTool) Description() string {
|
||||||
|
return "Read, write, list or search files, restricted to the configured data directory and explicit whitelist."
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parameters implements sdk.Tool.
|
||||||
|
func (f *fileTool) Parameters() (map[string]any, error) {
|
||||||
|
return map[string]any{
|
||||||
|
"type": "object",
|
||||||
|
"properties": map[string]any{
|
||||||
|
"action": map[string]any{"type": "string", "enum": []any{"read", "write", "list", "search"}, "description": "Operation to perform."},
|
||||||
|
"path": map[string]any{"type": "string", "description": "Absolute or data-directory-relative file path."},
|
||||||
|
"content": map[string]any{"type": "string", "description": "Body to write."},
|
||||||
|
"query": map[string]any{"type": "string", "description": "Substring filter for search."},
|
||||||
|
},
|
||||||
|
"required": []any{"action", "path"},
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Execute implements sdk.Tool.
|
||||||
|
func (f *fileTool) Execute(ctx context.Context, args any) (any, error) {
|
||||||
|
m, ok := args.(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("io: file: args must be an object")
|
||||||
|
}
|
||||||
|
switch action := argString(m, "action"); action {
|
||||||
|
case "read":
|
||||||
|
return f.read(ctx, argString(m, "path"))
|
||||||
|
case "write":
|
||||||
|
return f.write(ctx, argString(m, "path"), argString(m, "content"))
|
||||||
|
case "list":
|
||||||
|
return f.list(ctx, argString(m, "path"))
|
||||||
|
case "search":
|
||||||
|
return f.search(ctx, argString(m, "path"), argString(m, "query"))
|
||||||
|
default:
|
||||||
|
return nil, fmt.Errorf("io: file: unknown action %q", action)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// read returns a file's content, capped at maxRead bytes.
|
||||||
|
func (f *fileTool) read(ctx context.Context, p string) (any, error) {
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
real, err := f.resolveExisting(p)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
info, err := os.Stat(real)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: file: stat: %w", err)
|
||||||
|
}
|
||||||
|
if info.IsDir() {
|
||||||
|
return nil, fmt.Errorf("io: file: %q is a directory", real)
|
||||||
|
}
|
||||||
|
if info.Size() > f.maxRead {
|
||||||
|
return nil, fmt.Errorf("io: file: %q exceeds the %d-byte read limit", real, f.maxRead)
|
||||||
|
}
|
||||||
|
fh, err := os.Open(real)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: file: open: %w", err)
|
||||||
|
}
|
||||||
|
defer fh.Close()
|
||||||
|
data, err := stdio.ReadAll(stdio.LimitReader(fh, f.maxRead+1))
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: file: read: %w", err)
|
||||||
|
}
|
||||||
|
if int64(len(data)) > f.maxRead {
|
||||||
|
return nil, fmt.Errorf("io: file: %q exceeds the %d-byte read limit", real, f.maxRead)
|
||||||
|
}
|
||||||
|
return map[string]any{"path": real, "size": len(data), "content": string(data)}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// write creates or overwrites a file, capped at maxWrite bytes and confined to
|
||||||
|
// the allowed roots (parent dirs are created inside the root only).
|
||||||
|
func (f *fileTool) write(ctx context.Context, p, content string) (any, error) {
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
clean, err := f.resolve(p)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if int64(len(content)) > f.maxWrite {
|
||||||
|
return nil, fmt.Errorf("io: file: content exceeds the %d-byte write limit", f.maxWrite)
|
||||||
|
}
|
||||||
|
parent := filepath.Dir(clean)
|
||||||
|
if err := os.MkdirAll(parent, dirPerm); err != nil {
|
||||||
|
return nil, fmt.Errorf("io: file: mkdir: %w", err)
|
||||||
|
}
|
||||||
|
realParent, err := filepath.EvalSymlinks(parent)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: file: resolve parent: %w", err)
|
||||||
|
}
|
||||||
|
if !f.within(realParent) {
|
||||||
|
return nil, fmt.Errorf("io: file: parent directory escapes the allowed directories")
|
||||||
|
}
|
||||||
|
if info, err := os.Lstat(clean); err == nil && info.Mode()&os.ModeSymlink != 0 {
|
||||||
|
real, err := filepath.EvalSymlinks(clean)
|
||||||
|
if err != nil || !f.within(real) {
|
||||||
|
return nil, fmt.Errorf("io: file: %q escapes the allowed directories", p)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(clean, []byte(content), filePerm); err != nil {
|
||||||
|
return nil, fmt.Errorf("io: file: write: %w", err)
|
||||||
|
}
|
||||||
|
return map[string]any{"path": clean, "size": len(content)}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// list returns the entries of a directory, capped at maxList.
|
||||||
|
func (f *fileTool) list(ctx context.Context, p string) (any, error) {
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
real, err := f.resolveExisting(p)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
info, err := os.Stat(real)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: file: stat: %w", err)
|
||||||
|
}
|
||||||
|
if !info.IsDir() {
|
||||||
|
return nil, fmt.Errorf("io: file: %q is not a directory", real)
|
||||||
|
}
|
||||||
|
entries, err := os.ReadDir(real)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: file: readdir: %w", err)
|
||||||
|
}
|
||||||
|
if len(entries) > f.maxList {
|
||||||
|
entries = entries[:f.maxList]
|
||||||
|
}
|
||||||
|
out := make([]map[string]any, 0, len(entries))
|
||||||
|
for _, e := range entries {
|
||||||
|
ei, err := e.Info()
|
||||||
|
if err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
out = append(out, map[string]any{"name": e.Name(), "size": ei.Size(), "dir": ei.IsDir()})
|
||||||
|
}
|
||||||
|
return map[string]any{"path": real, "entries": out, "count": len(out)}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// search walks a base directory (default the primary root) and returns files
|
||||||
|
// whose names contain query, capped at maxList results.
|
||||||
|
func (f *fileTool) search(ctx context.Context, base, query string) (any, error) {
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(query) == "" {
|
||||||
|
return nil, fmt.Errorf("io: file: search query is required")
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(base) == "" {
|
||||||
|
if len(f.roots) == 0 {
|
||||||
|
return nil, fmt.Errorf("io: file: no base directory configured")
|
||||||
|
}
|
||||||
|
base = f.roots[0]
|
||||||
|
}
|
||||||
|
real, err := f.resolveExisting(base)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
info, err := os.Stat(real)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: file: stat: %w", err)
|
||||||
|
}
|
||||||
|
if !info.IsDir() {
|
||||||
|
return nil, fmt.Errorf("io: file: %q is not a directory", real)
|
||||||
|
}
|
||||||
|
q := strings.ToLower(query)
|
||||||
|
hits := make([]map[string]any, 0)
|
||||||
|
walkErr := filepath.WalkDir(real, func(pp string, d fs.DirEntry, err error) error {
|
||||||
|
if err != nil || d.IsDir() {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if !strings.Contains(strings.ToLower(d.Name()), q) {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
di, err := d.Info()
|
||||||
|
if err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
rel, rerr := filepath.Rel(real, pp)
|
||||||
|
if rerr != nil {
|
||||||
|
rel = pp
|
||||||
|
}
|
||||||
|
hits = append(hits, map[string]any{"path": rel, "size": di.Size()})
|
||||||
|
if len(hits) >= f.maxList {
|
||||||
|
return fs.SkipAll
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
if walkErr != nil && !errors.Is(walkErr, fs.SkipAll) {
|
||||||
|
return nil, fmt.Errorf("io: file: walk: %w", walkErr)
|
||||||
|
}
|
||||||
|
return map[string]any{"base": real, "query": query, "hits": hits, "count": len(hits)}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// resolve cleans p and verifies it stays within one of the allowed roots. It
|
||||||
|
// rejects empty paths and any path that escapes via ".." or an absolute target
|
||||||
|
// outside the roots.
|
||||||
|
func (f *fileTool) resolve(p string) (string, error) {
|
||||||
|
if strings.TrimSpace(p) == "" {
|
||||||
|
return "", fmt.Errorf("io: file: path is required")
|
||||||
|
}
|
||||||
|
if len(f.roots) == 0 {
|
||||||
|
return "", fmt.Errorf("io: file: no base directory configured")
|
||||||
|
}
|
||||||
|
clean := filepath.Clean(p)
|
||||||
|
if !filepath.IsAbs(clean) {
|
||||||
|
clean = filepath.Join(f.roots[0], clean)
|
||||||
|
}
|
||||||
|
clean = filepath.Clean(clean)
|
||||||
|
if !f.within(clean) {
|
||||||
|
return "", fmt.Errorf("io: file: %q is outside the allowed directories", p)
|
||||||
|
}
|
||||||
|
return clean, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// resolveExisting resolves p and verifies the symlink-resolved real path stays
|
||||||
|
// within the roots, guarding against escaping through a symlink. It errors when
|
||||||
|
// the target does not exist.
|
||||||
|
func (f *fileTool) resolveExisting(p string) (string, error) {
|
||||||
|
clean, err := f.resolve(p)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
real, err := filepath.EvalSymlinks(clean)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("io: file: resolve %q: %w", p, err)
|
||||||
|
}
|
||||||
|
if !f.within(real) {
|
||||||
|
f.log.Warn("io file: rejected path escaping allowed directories", "path", p, "resolved", real)
|
||||||
|
return "", fmt.Errorf("io: file: %q escapes the allowed directories", p)
|
||||||
|
}
|
||||||
|
return real, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// within reports whether p is inside (or equal to) one of the allowed roots.
|
||||||
|
func (f *fileTool) within(p string) bool {
|
||||||
|
for _, root := range f.roots {
|
||||||
|
rel, err := filepath.Rel(root, p)
|
||||||
|
if err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if rel == "." || (rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator))) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// absoluteClean absolutizes and cleans p, returning the best-effort absolute
|
||||||
|
// form used as an allowed root.
|
||||||
|
func absoluteClean(p string) string {
|
||||||
|
a, err := filepath.Abs(p)
|
||||||
|
if err != nil {
|
||||||
|
return filepath.Clean(p)
|
||||||
|
}
|
||||||
|
return filepath.Clean(a)
|
||||||
|
}
|
||||||
@@ -0,0 +1,187 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
// fileToolForTest builds a tools set over cfg with opts and returns the file tool.
|
||||||
|
func fileToolForTest(t *testing.T, cfg *config.Config, opts ...Option) *fileTool {
|
||||||
|
t.Helper()
|
||||||
|
return findTool(t, buildTools(t, cfg, opts...), "file").(*fileTool)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileTool_WriteAndRead(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
dir := t.TempDir()
|
||||||
|
ft := fileToolForTest(t, &config.Config{DataDir: dir})
|
||||||
|
|
||||||
|
if _, err := ft.Execute(context.Background(), map[string]any{"action": "write", "path": "sub/a.txt", "content": "hello world"}); err != nil {
|
||||||
|
t.Fatalf("write: %v", err)
|
||||||
|
}
|
||||||
|
got, err := ft.Execute(context.Background(), map[string]any{"action": "read", "path": "sub/a.txt"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("read: %v", err)
|
||||||
|
}
|
||||||
|
m := got.(map[string]any)
|
||||||
|
if m["content"] != "hello world" {
|
||||||
|
t.Fatalf("content = %q", m["content"])
|
||||||
|
}
|
||||||
|
if m["size"] != len("hello world") {
|
||||||
|
t.Fatalf("size = %v", m["size"])
|
||||||
|
}
|
||||||
|
if m["path"] != filepath.Join(dir, "sub", "a.txt") {
|
||||||
|
t.Fatalf("path = %q", m["path"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileTool_RejectsTraversal(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ft := fileToolForTest(t, &config.Config{DataDir: t.TempDir()})
|
||||||
|
cases := []string{
|
||||||
|
"../escape.txt",
|
||||||
|
"../../../../etc/passwd",
|
||||||
|
"sub/../../escape.txt",
|
||||||
|
}
|
||||||
|
for _, p := range cases {
|
||||||
|
_, err := ft.Execute(context.Background(), map[string]any{"action": "read", "path": p})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatalf("expected traversal rejection for %q", p)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// An absolute path outside the data dir is also rejected.
|
||||||
|
outside := filepath.Join(t.TempDir(), "outside.txt")
|
||||||
|
_, err := ft.Execute(context.Background(), map[string]any{"action": "write", "path": outside, "content": "x"})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected rejection for absolute path outside the data dir")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileTool_ReadSizeLimit(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
dir := t.TempDir()
|
||||||
|
// Write with the default limit, then read with a tiny cap.
|
||||||
|
if _, err := fileToolForTest(t, &config.Config{DataDir: dir}).
|
||||||
|
Execute(context.Background(), map[string]any{"action": "write", "path": "big.txt", "content": "1234567890"}); err != nil {
|
||||||
|
t.Fatalf("write: %v", err)
|
||||||
|
}
|
||||||
|
ft := fileToolForTest(t, &config.Config{DataDir: dir}, WithReadLimit(4))
|
||||||
|
_, err := ft.Execute(context.Background(), map[string]any{"action": "read", "path": "big.txt"})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "read limit") {
|
||||||
|
t.Fatalf("expected read-limit error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileTool_WriteSizeLimit(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ft := fileToolForTest(t, &config.Config{DataDir: t.TempDir()}, WithWriteLimit(4))
|
||||||
|
_, err := ft.Execute(context.Background(), map[string]any{"action": "write", "path": "x.txt", "content": "too long"})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "write limit") {
|
||||||
|
t.Fatalf("expected write-limit error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileTool_List(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
dir := t.TempDir()
|
||||||
|
ft := fileToolForTest(t, &config.Config{DataDir: dir})
|
||||||
|
for _, name := range []string{"a.txt", "b.txt", "sub"} {
|
||||||
|
p := filepath.Join(dir, name)
|
||||||
|
if name == "sub" {
|
||||||
|
if err := os.MkdirAll(p, dirPerm); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
} else if err := os.WriteFile(p, []byte("x"), filePerm); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
got, err := ft.Execute(context.Background(), map[string]any{"action": "list", "path": "."})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("list: %v", err)
|
||||||
|
}
|
||||||
|
m := got.(map[string]any)
|
||||||
|
if m["count"].(int) != 3 {
|
||||||
|
t.Fatalf("count = %v", m["count"])
|
||||||
|
}
|
||||||
|
entries := m["entries"].([]map[string]any)
|
||||||
|
found := map[string]bool{}
|
||||||
|
for _, e := range entries {
|
||||||
|
found[e["name"].(string)] = true
|
||||||
|
}
|
||||||
|
for _, want := range []string{"a.txt", "b.txt", "sub"} {
|
||||||
|
if !found[want] {
|
||||||
|
t.Errorf("missing entry %q", want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileTool_Search(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
dir := t.TempDir()
|
||||||
|
ft := fileToolForTest(t, &config.Config{DataDir: dir})
|
||||||
|
for _, name := range []string{"report_2026.txt", "notes.md", "report_old.txt"} {
|
||||||
|
if err := os.WriteFile(filepath.Join(dir, name), []byte("x"), filePerm); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
got, err := ft.Execute(context.Background(), map[string]any{"action": "search", "query": "report"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("search: %v", err)
|
||||||
|
}
|
||||||
|
m := got.(map[string]any)
|
||||||
|
if m["count"].(int) != 2 {
|
||||||
|
t.Fatalf("count = %v", m["count"])
|
||||||
|
}
|
||||||
|
if _, err := ft.Execute(context.Background(), map[string]any{"action": "search", "query": ""}); err == nil {
|
||||||
|
t.Fatal("expected error for empty search query")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileTool_NoBaseDirectory(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ft := fileToolForTest(t, &config.Config{}) // empty DataDir, no whitelist
|
||||||
|
_, err := ft.Execute(context.Background(), map[string]any{"action": "read", "path": "a.txt"})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "no base directory") {
|
||||||
|
t.Fatalf("expected no-base-directory error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileTool_PathWhitelist(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
dir := t.TempDir()
|
||||||
|
extra := t.TempDir()
|
||||||
|
ft := fileToolForTest(t, &config.Config{DataDir: dir}, WithFilePathWhitelist(extra))
|
||||||
|
if err := os.WriteFile(filepath.Join(extra, "w.txt"), []byte("hi"), filePerm); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
got, err := ft.Execute(context.Background(), map[string]any{"action": "read", "path": filepath.Join(extra, "w.txt")})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("read whitelisted file: %v", err)
|
||||||
|
}
|
||||||
|
if got.(map[string]any)["content"] != "hi" {
|
||||||
|
t.Fatalf("content = %q", got.(map[string]any)["content"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileTool_RejectsSymlinkEscape(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
dir := t.TempDir()
|
||||||
|
secret := t.TempDir()
|
||||||
|
if err := os.WriteFile(filepath.Join(secret, "secret.txt"), []byte("hidden"), filePerm); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
link := filepath.Join(dir, "link.txt")
|
||||||
|
if err := os.Symlink(filepath.Join(secret, "secret.txt"), link); err != nil {
|
||||||
|
t.Skipf("symlink unsupported: %v", err)
|
||||||
|
}
|
||||||
|
ft := fileToolForTest(t, &config.Config{DataDir: dir})
|
||||||
|
_, err := ft.Execute(context.Background(), map[string]any{"action": "read", "path": "link.txt"})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected symlink escape to be rejected")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,54 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"io"
|
||||||
|
"log/slog"
|
||||||
|
"net/http/httptest"
|
||||||
|
"net/url"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||||
|
)
|
||||||
|
|
||||||
|
// silentLogger is the base logger used by test fixtures so that expected
|
||||||
|
// security-denial WARN records are not dumped into the test output.
|
||||||
|
func silentLogger() *slog.Logger {
|
||||||
|
return logger.New(logger.Options{Out: io.Discard})
|
||||||
|
}
|
||||||
|
|
||||||
|
// buildTools builds a tool set over cfg with a silent logger, applying opts on
|
||||||
|
// top (caller options win over the default discard logger).
|
||||||
|
func buildTools(t *testing.T, cfg *config.Config, opts ...Option) []sdk.Tool {
|
||||||
|
t.Helper()
|
||||||
|
all := append([]Option{WithLogger(silentLogger())}, opts...)
|
||||||
|
set, err := New(cfg, all...)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
return set
|
||||||
|
}
|
||||||
|
|
||||||
|
// findTool returns the tool with the given name from a set, failing the test
|
||||||
|
// when it is absent.
|
||||||
|
func findTool(t *testing.T, set []sdk.Tool, name string) sdk.Tool {
|
||||||
|
t.Helper()
|
||||||
|
for _, tool := range set {
|
||||||
|
if tool.Name() == name {
|
||||||
|
return tool
|
||||||
|
}
|
||||||
|
}
|
||||||
|
t.Fatalf("tool %q not found in set", name)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// serverHost returns the bare host (no port, no brackets) of a test server, the
|
||||||
|
// value to put on the SSRF allowlist so the tools may reach it.
|
||||||
|
func serverHost(srv *httptest.Server) string {
|
||||||
|
u, err := url.Parse(srv.URL)
|
||||||
|
if err != nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return u.Hostname()
|
||||||
|
}
|
||||||
@@ -0,0 +1,127 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
stdio "io"
|
||||||
|
"log/slog"
|
||||||
|
"mime"
|
||||||
|
"net/http"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Compile-time assertion that httpTool implements sdk.Tool.
|
||||||
|
var _ sdk.Tool = (*httpTool)(nil)
|
||||||
|
|
||||||
|
// httpTool performs a single GET or POST request with headers, body, a timeout
|
||||||
|
// and SSRF protection, truncating the response to a size cap.
|
||||||
|
type httpTool struct {
|
||||||
|
client *http.Client
|
||||||
|
allowedHosts []string
|
||||||
|
maxBytes int64
|
||||||
|
log *slog.Logger
|
||||||
|
}
|
||||||
|
|
||||||
|
func newHTTPTool(t *tools) *httpTool {
|
||||||
|
return &httpTool{
|
||||||
|
client: t.client,
|
||||||
|
allowedHosts: t.allowedHosts,
|
||||||
|
maxBytes: t.maxHTTPBytes,
|
||||||
|
log: t.log,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Name implements sdk.Tool.
|
||||||
|
func (h *httpTool) Name() string { return "http" }
|
||||||
|
|
||||||
|
// Description implements sdk.Tool.
|
||||||
|
func (h *httpTool) Description() string {
|
||||||
|
return "Perform an HTTP GET or POST request with headers, body, timeout and SSRF protection."
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parameters implements sdk.Tool.
|
||||||
|
func (h *httpTool) Parameters() (map[string]any, error) {
|
||||||
|
return map[string]any{
|
||||||
|
"type": "object",
|
||||||
|
"properties": map[string]any{
|
||||||
|
"method": map[string]any{"type": "string", "enum": []any{"GET", "POST"}, "description": "HTTP method."},
|
||||||
|
"url": map[string]any{"type": "string", "description": "Absolute http(s) URL."},
|
||||||
|
"headers": map[string]any{"type": "object", "description": "Request headers."},
|
||||||
|
"body": map[string]any{"type": "string", "description": "Request body for POST."},
|
||||||
|
"timeout_ms": map[string]any{"type": "integer", "description": "Per-request timeout in milliseconds."},
|
||||||
|
"max_bytes": map[string]any{"type": "integer", "description": "Response size cap in bytes."},
|
||||||
|
},
|
||||||
|
"required": []any{"url"},
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Execute implements sdk.Tool.
|
||||||
|
func (h *httpTool) Execute(ctx context.Context, args any) (any, error) {
|
||||||
|
m, ok := args.(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("io: http: args must be an object")
|
||||||
|
}
|
||||||
|
method := strings.ToUpper(argString(m, "method"))
|
||||||
|
if method == "" {
|
||||||
|
method = http.MethodGet
|
||||||
|
}
|
||||||
|
if method != http.MethodGet && method != http.MethodPost {
|
||||||
|
return nil, fmt.Errorf("io: http: unsupported method %q", method)
|
||||||
|
}
|
||||||
|
rawURL := argString(m, "url")
|
||||||
|
if strings.TrimSpace(rawURL) == "" {
|
||||||
|
return nil, fmt.Errorf("io: http: url is required")
|
||||||
|
}
|
||||||
|
u, err := validateRequestURL(ctx, rawURL, h.allowedHosts)
|
||||||
|
if err != nil {
|
||||||
|
h.log.Warn("io http: blocked request", "url", rawURL, "err", err)
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
body := argString(m, "body")
|
||||||
|
var reqBody stdio.Reader
|
||||||
|
if body != "" {
|
||||||
|
reqBody = strings.NewReader(body)
|
||||||
|
}
|
||||||
|
req, err := http.NewRequestWithContext(ctx, method, u.String(), reqBody)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: http: build request: %w", err)
|
||||||
|
}
|
||||||
|
for k, v := range argStringMap(m, "headers") {
|
||||||
|
req.Header.Set(k, v)
|
||||||
|
}
|
||||||
|
ct := req.Header.Get("Content-Type")
|
||||||
|
if ct != "" {
|
||||||
|
if _, err := validateContentType(ct); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
} else if body != "" {
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
}
|
||||||
|
if ms, ok := argInt(m, "timeout_ms"); ok && ms > 0 {
|
||||||
|
ctx, cancel := context.WithTimeout(ctx, time.Duration(ms)*time.Millisecond)
|
||||||
|
defer cancel()
|
||||||
|
req = req.WithContext(ctx)
|
||||||
|
}
|
||||||
|
if h.client == nil {
|
||||||
|
return nil, fmt.Errorf("io: http: no underlying client")
|
||||||
|
}
|
||||||
|
resp, err := h.client.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: http: do: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
data, truncated, err := readLimited(resp.Body, h.maxBytes)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: http: read: %w", err)
|
||||||
|
}
|
||||||
|
ctOut, _, _ := mime.ParseMediaType(resp.Header.Get("Content-Type"))
|
||||||
|
return map[string]any{
|
||||||
|
"status": resp.StatusCode,
|
||||||
|
"content_type": ctOut,
|
||||||
|
"body": string(data),
|
||||||
|
"truncated": truncated,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,189 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
// httpToolForTest builds a tools set with opts and returns the http tool.
|
||||||
|
func httpToolForTest(t *testing.T, opts ...Option) *httpTool {
|
||||||
|
t.Helper()
|
||||||
|
return findTool(t, buildTools(t, &config.Config{}, opts...), "http").(*httpTool)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHTTPTool_Get(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, _ = io.WriteString(w, `{"ok":true,"path":"`+r.URL.Path+`"}`)
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||||
|
|
||||||
|
got, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": srv.URL + "/hello"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Execute: %v", err)
|
||||||
|
}
|
||||||
|
m := got.(map[string]any)
|
||||||
|
if m["status"] != http.StatusOK {
|
||||||
|
t.Fatalf("status = %v", m["status"])
|
||||||
|
}
|
||||||
|
body := m["body"].(string)
|
||||||
|
if !strings.Contains(body, `"ok":true`) || !strings.Contains(body, "/hello") {
|
||||||
|
t.Fatalf("body = %q", body)
|
||||||
|
}
|
||||||
|
if m["truncated"].(bool) {
|
||||||
|
t.Fatal("unexpected truncation")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHTTPTool_Post(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method != http.MethodPost {
|
||||||
|
t.Fatalf("method = %s", r.Method)
|
||||||
|
}
|
||||||
|
body, _ := io.ReadAll(r.Body)
|
||||||
|
w.Header().Set("Content-Type", "text/plain")
|
||||||
|
_, _ = w.Write(body)
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||||
|
|
||||||
|
got, err := ht.Execute(context.Background(), map[string]any{
|
||||||
|
"method": "POST",
|
||||||
|
"url": srv.URL,
|
||||||
|
"body": "name=cyrene",
|
||||||
|
"headers": map[string]any{"Content-Type": "application/x-www-form-urlencoded"},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Execute: %v", err)
|
||||||
|
}
|
||||||
|
if body := got.(map[string]any)["body"].(string); body != "name=cyrene" {
|
||||||
|
t.Fatalf("echoed body = %q", body)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHTTPTool_UnsupportedMethod(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ht := httpToolForTest(t)
|
||||||
|
_, err := ht.Execute(context.Background(), map[string]any{"method": "DELETE", "url": "http://example.com"})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "unsupported method") {
|
||||||
|
t.Fatalf("expected unsupported-method error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHTTPTool_Timeout(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
time.Sleep(200 * time.Millisecond)
|
||||||
|
_, _ = w.Write([]byte("late"))
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPTimeout(30*time.Millisecond))
|
||||||
|
|
||||||
|
_, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": srv.URL})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected timeout error")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHTTPTool_RejectsBadContentType(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ht := httpToolForTest(t)
|
||||||
|
_, err := ht.Execute(context.Background(), map[string]any{
|
||||||
|
"method": "POST",
|
||||||
|
"url": "http://93.184.216.34/",
|
||||||
|
"body": "x",
|
||||||
|
"headers": map[string]any{"Content-Type": "not a real content type {{"},
|
||||||
|
})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "invalid content-type") {
|
||||||
|
t.Fatalf("expected invalid content-type error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHTTPTool_ResponseTruncation(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
payload := strings.Repeat("a", 5000)
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
_, _ = io.WriteString(w, payload)
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()), WithHTTPResponseLimit(1000))
|
||||||
|
|
||||||
|
got, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": srv.URL})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Execute: %v", err)
|
||||||
|
}
|
||||||
|
m := got.(map[string]any)
|
||||||
|
if !m["truncated"].(bool) {
|
||||||
|
t.Fatal("expected truncated=true")
|
||||||
|
}
|
||||||
|
if body := m["body"].(string); len(body) != 1000 {
|
||||||
|
t.Fatalf("body len = %d", len(body))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHTTPTool_SSRFDenied(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ht := httpToolForTest(t) // no whitelist
|
||||||
|
cases := []string{
|
||||||
|
"http://127.0.0.1:9/",
|
||||||
|
"http://localhost:9/",
|
||||||
|
"http://10.0.0.1/",
|
||||||
|
"http://172.16.0.1/",
|
||||||
|
"http://192.168.1.1/",
|
||||||
|
"http://169.254.1.1/",
|
||||||
|
"http://0.0.0.0/",
|
||||||
|
"http://[::1]:9/",
|
||||||
|
}
|
||||||
|
for _, u := range cases {
|
||||||
|
_, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": u})
|
||||||
|
if !errors.Is(err, ErrSSRF) {
|
||||||
|
t.Fatalf("url %q: expected ErrSSRF, got %v", u, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHTTPTool_SSRFAllowedViaWhitelist(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
_, _ = io.WriteString(w, "ok")
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
ht := httpToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||||
|
|
||||||
|
got, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": srv.URL})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Execute (whitelisted): %v", err)
|
||||||
|
}
|
||||||
|
if body := got.(map[string]any)["body"].(string); body != "ok" {
|
||||||
|
t.Fatalf("body = %q", body)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHTTPTool_UnsupportedScheme(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ht := httpToolForTest(t)
|
||||||
|
_, err := ht.Execute(context.Background(), map[string]any{"method": "GET", "url": "ftp://example.com/x"})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "unsupported scheme") {
|
||||||
|
t.Fatalf("expected unsupported-scheme error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHTTPTool_MissingURL(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ht := httpToolForTest(t)
|
||||||
|
_, err := ht.Execute(context.Background(), map[string]any{"method": "GET"})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "url is required") {
|
||||||
|
t.Fatalf("expected url-required error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,101 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory"
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||||
|
)
|
||||||
|
|
||||||
|
// KnowledgeRetriever is the subset of the runtime memory retriever the knowledge
|
||||||
|
// tool depends on. *memory.Retriever implements it; nil means no knowledge base
|
||||||
|
// is configured.
|
||||||
|
type KnowledgeRetriever interface {
|
||||||
|
Retrieve(ctx context.Context, personaID, query string, topK int) ([]memory.Result, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Compile-time assertion that knowledgeTool implements sdk.Tool.
|
||||||
|
var _ sdk.Tool = (*knowledgeTool)(nil)
|
||||||
|
|
||||||
|
// knowledgeTool answers a question from the configured knowledge base via a
|
||||||
|
// memory retriever, returning ranked results or a clear message when absent.
|
||||||
|
type knowledgeTool struct {
|
||||||
|
retriever KnowledgeRetriever
|
||||||
|
maxResults int
|
||||||
|
log *slog.Logger
|
||||||
|
}
|
||||||
|
|
||||||
|
func newKnowledgeTool(t *tools) *knowledgeTool {
|
||||||
|
return &knowledgeTool{
|
||||||
|
retriever: t.retriever,
|
||||||
|
maxResults: t.maxResults,
|
||||||
|
log: t.log,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Name implements sdk.Tool.
|
||||||
|
func (k *knowledgeTool) Name() string { return "knowledge" }
|
||||||
|
|
||||||
|
// Description implements sdk.Tool.
|
||||||
|
func (k *knowledgeTool) Description() string {
|
||||||
|
return "Query the knowledge base and long-term memory store for context relevant to a question."
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parameters implements sdk.Tool.
|
||||||
|
func (k *knowledgeTool) Parameters() (map[string]any, error) {
|
||||||
|
return map[string]any{
|
||||||
|
"type": "object",
|
||||||
|
"properties": map[string]any{
|
||||||
|
"query": map[string]any{"type": "string", "description": "Query topic."},
|
||||||
|
"persona": map[string]any{"type": "string", "description": "Persona scope id (optional)."},
|
||||||
|
"top_k": map[string]any{"type": "integer", "description": "Max results."},
|
||||||
|
},
|
||||||
|
"required": []any{"query"},
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Execute implements sdk.Tool.
|
||||||
|
func (k *knowledgeTool) Execute(ctx context.Context, args any) (any, error) {
|
||||||
|
if k.retriever == nil {
|
||||||
|
return nil, fmt.Errorf("io: knowledge: no knowledge base configured")
|
||||||
|
}
|
||||||
|
m, ok := args.(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("io: knowledge: args must be an object")
|
||||||
|
}
|
||||||
|
query := argString(m, "query")
|
||||||
|
if strings.TrimSpace(query) == "" {
|
||||||
|
return nil, fmt.Errorf("io: knowledge: query is required")
|
||||||
|
}
|
||||||
|
persona := argString(m, "persona")
|
||||||
|
topK := k.maxResults
|
||||||
|
if v, ok := argInt(m, "top_k"); ok && v > 0 {
|
||||||
|
topK = v
|
||||||
|
}
|
||||||
|
results, err := k.retriever.Retrieve(ctx, persona, query, topK)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: knowledge: retrieve: %w", err)
|
||||||
|
}
|
||||||
|
out := make([]map[string]any, 0, len(results))
|
||||||
|
for _, r := range results {
|
||||||
|
out = append(out, map[string]any{
|
||||||
|
"id": r.Item.ID,
|
||||||
|
"kind": string(r.Item.Kind),
|
||||||
|
"scope": string(r.Item.Scope),
|
||||||
|
"content": r.Item.Content,
|
||||||
|
"score": r.Score,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
if len(out) == 0 {
|
||||||
|
return map[string]any{
|
||||||
|
"query": query,
|
||||||
|
"count": 0,
|
||||||
|
"note": "no relevant knowledge found",
|
||||||
|
"results": []map[string]any{},
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
return map[string]any{"query": query, "count": len(out), "results": out}, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,99 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory"
|
||||||
|
)
|
||||||
|
|
||||||
|
// stubRetriever is a fake KnowledgeRetriever for testing knowledge responses.
|
||||||
|
type stubRetriever struct {
|
||||||
|
results []memory.Result
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *stubRetriever) Retrieve(ctx context.Context, personaID, query string, topK int) ([]memory.Result, error) {
|
||||||
|
_ = ctx
|
||||||
|
return s.results, s.err
|
||||||
|
}
|
||||||
|
|
||||||
|
// knowledgeToolForTest builds a tools set with opts and returns the knowledge
|
||||||
|
// tool.
|
||||||
|
func knowledgeToolForTest(t *testing.T, opts ...Option) *knowledgeTool {
|
||||||
|
t.Helper()
|
||||||
|
return findTool(t, buildTools(t, &config.Config{}, opts...), "knowledge").(*knowledgeTool)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestKnowledge_NoKnowledgeBase(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
kt := knowledgeToolForTest(t) // no retriever wired
|
||||||
|
_, err := kt.Execute(context.Background(), map[string]any{"query": "cyrene"})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "no knowledge base") {
|
||||||
|
t.Fatalf("expected no-knowledge-base error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestKnowledge_ReturnsResults(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
sr := &stubRetriever{results: []memory.Result{
|
||||||
|
{Item: memory.MemoryItem{ID: "1", Content: "Cyrene 是主人", Kind: memory.KindWorldEvent, Scope: memory.ScopeProject}, Score: 0.9},
|
||||||
|
}}
|
||||||
|
kt := knowledgeToolForTest(t, WithKnowledgeRetriever(sr))
|
||||||
|
|
||||||
|
got, err := kt.Execute(context.Background(), map[string]any{"query": "cyrene", "persona": "p1", "top_k": 5})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Execute: %v", err)
|
||||||
|
}
|
||||||
|
m := got.(map[string]any)
|
||||||
|
if m["count"].(int) != 1 {
|
||||||
|
t.Fatalf("count = %v", m["count"])
|
||||||
|
}
|
||||||
|
results := m["results"].([]map[string]any)
|
||||||
|
if results[0]["content"] != "Cyrene 是主人" {
|
||||||
|
t.Fatalf("content = %v", results[0]["content"])
|
||||||
|
}
|
||||||
|
if results[0]["kind"] != "world_event" {
|
||||||
|
t.Fatalf("kind = %v", results[0]["kind"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestKnowledge_EmptyResults(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
kt := knowledgeToolForTest(t, WithKnowledgeRetriever(&stubRetriever{}))
|
||||||
|
got, err := kt.Execute(context.Background(), map[string]any{"query": "nothing"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Execute: %v", err)
|
||||||
|
}
|
||||||
|
m := got.(map[string]any)
|
||||||
|
if m["count"].(int) != 0 {
|
||||||
|
t.Fatalf("count = %v", m["count"])
|
||||||
|
}
|
||||||
|
if m["note"] != "no relevant knowledge found" {
|
||||||
|
t.Fatalf("note = %v", m["note"])
|
||||||
|
}
|
||||||
|
if _, ok := m["results"].([]map[string]any); !ok {
|
||||||
|
t.Fatalf("results absent or wrong type: %T", m["results"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestKnowledge_RetrieverError(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
kt := knowledgeToolForTest(t, WithKnowledgeRetriever(&stubRetriever{err: errors.New("boom")}))
|
||||||
|
_, err := kt.Execute(context.Background(), map[string]any{"query": "x"})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "retrieve") {
|
||||||
|
t.Fatalf("expected retrieve error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestKnowledge_MissingQuery(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
kt := knowledgeToolForTest(t, WithKnowledgeRetriever(&stubRetriever{}))
|
||||||
|
_, err := kt.Execute(context.Background(), map[string]any{})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "query is required") {
|
||||||
|
t.Fatalf("expected query-required error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,101 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"net"
|
||||||
|
"net/url"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ErrSSRF is returned when a requested URL targets a loopback, private or
|
||||||
|
// link-local address or a localhost hostname that is not on the allowlist.
|
||||||
|
// Use errors.Is to test for it.
|
||||||
|
var ErrSSRF = errors.New("io: disallowed url")
|
||||||
|
|
||||||
|
// validateRequestURL parses rawURL and refuses SSRF-prone targets: it returns an
|
||||||
|
// error wrapping ErrSSRF for loopback/private/link-local/unspecified addresses
|
||||||
|
// and for localhost hostnames, unless the host (or hostname) is on the explicit
|
||||||
|
// allowlist. A hostname that resolves to any denied address is also refused.
|
||||||
|
func validateRequestURL(ctx context.Context, rawURL string, allowed []string) (*url.URL, error) {
|
||||||
|
if strings.TrimSpace(rawURL) == "" {
|
||||||
|
return nil, fmt.Errorf("io: url is required")
|
||||||
|
}
|
||||||
|
u, err := url.Parse(rawURL)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: parse url: %w", err)
|
||||||
|
}
|
||||||
|
if u.Scheme != "http" && u.Scheme != "https" {
|
||||||
|
return nil, fmt.Errorf("io: unsupported scheme %q", u.Scheme)
|
||||||
|
}
|
||||||
|
host := hostOnly(u.Host)
|
||||||
|
if host == "" {
|
||||||
|
return nil, fmt.Errorf("io: url has no host")
|
||||||
|
}
|
||||||
|
if isHostAllowed(host, allowed) {
|
||||||
|
return u, nil
|
||||||
|
}
|
||||||
|
if ip := net.ParseIP(host); ip != nil {
|
||||||
|
if isDeniedIP(ip) {
|
||||||
|
return nil, fmt.Errorf("%w: %s", ErrSSRF, host)
|
||||||
|
}
|
||||||
|
return u, nil
|
||||||
|
}
|
||||||
|
addrs, err := net.DefaultResolver.LookupHost(ctx, host)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: resolve %q: %w", host, err)
|
||||||
|
}
|
||||||
|
for _, a := range addrs {
|
||||||
|
if ip := net.ParseIP(a); ip != nil && isDeniedIP(ip) {
|
||||||
|
return nil, fmt.Errorf("%w: %q resolves to %q", ErrSSRF, host, a)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return u, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// isDeniedIP reports whether ip is a loopback, private, link-local, multicast or
|
||||||
|
// unspecified address — the targets SSRF protection refuses by default.
|
||||||
|
func isDeniedIP(ip net.IP) bool {
|
||||||
|
if ip == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return ip.IsLoopback() ||
|
||||||
|
ip.IsPrivate() ||
|
||||||
|
ip.IsLinkLocalUnicast() ||
|
||||||
|
ip.IsLinkLocalMulticast() ||
|
||||||
|
ip.IsUnspecified()
|
||||||
|
}
|
||||||
|
|
||||||
|
// isHostAllowed reports whether the host (with any port stripped) matches an
|
||||||
|
// entry on the allowlist.
|
||||||
|
func isHostAllowed(host string, allowed []string) bool {
|
||||||
|
h := strings.ToLower(hostOnly(host))
|
||||||
|
if h == "" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
for _, a := range allowed {
|
||||||
|
if strings.EqualFold(hostOnly(a), h) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// hostOnly strips the port (and any IPv6 brackets) from a host:port string.
|
||||||
|
func hostOnly(hostport string) string {
|
||||||
|
h := strings.TrimSpace(hostport)
|
||||||
|
if h == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
if strings.HasPrefix(h, "[") {
|
||||||
|
if j := strings.Index(h, "]"); j >= 0 {
|
||||||
|
return h[1:j]
|
||||||
|
}
|
||||||
|
return h
|
||||||
|
}
|
||||||
|
if i := strings.LastIndex(h, ":"); i >= 0 {
|
||||||
|
return h[:i]
|
||||||
|
}
|
||||||
|
return h
|
||||||
|
}
|
||||||
@@ -0,0 +1,232 @@
|
|||||||
|
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 {
|
||||||
|
return []sdk.Tool{
|
||||||
|
newFileTool(t),
|
||||||
|
newHTTPTool(t),
|
||||||
|
newWebSearchTool(t),
|
||||||
|
newWebFetchTool(t),
|
||||||
|
newKnowledgeTool(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
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,107 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestTools_ReturnsAllFiveTools(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
set, err := New(&config.Config{DataDir: t.TempDir()})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
if len(set) != 5 {
|
||||||
|
t.Fatalf("expected 5 tools, got %d", len(set))
|
||||||
|
}
|
||||||
|
want := []string{"file", "http", "web-search", "web-fetch", "knowledge"}
|
||||||
|
seen := map[string]bool{}
|
||||||
|
for _, tool := range set {
|
||||||
|
seen[tool.Name()] = true
|
||||||
|
if tool.Description() == "" {
|
||||||
|
t.Errorf("tool %q has an empty description", tool.Name())
|
||||||
|
}
|
||||||
|
params, err := tool.Parameters()
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("tool %q Parameters error: %v", tool.Name(), err)
|
||||||
|
}
|
||||||
|
if params == nil {
|
||||||
|
t.Errorf("tool %q has nil parameters", tool.Name())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, name := range want {
|
||||||
|
if !seen[name] {
|
||||||
|
t.Errorf("missing tool %q", name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNew_NilConfigErrors(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
if _, err := New(nil); err == nil {
|
||||||
|
t.Fatal("expected error for nil config")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRegister_RegistersEveryTool(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
reg := registry.New()
|
||||||
|
if err := Register(reg); err != nil {
|
||||||
|
t.Fatalf("Register: %v", err)
|
||||||
|
}
|
||||||
|
for _, name := range []string{"file", "http", "web-search", "web-fetch", "knowledge"} {
|
||||||
|
c, err := reg.Get(name)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Get(%q): %v", name, err)
|
||||||
|
}
|
||||||
|
if c.Kind != registry.KindTool {
|
||||||
|
t.Errorf("kind of %q = %q, want tool", name, c.Kind)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRegister_NilRegistrar(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
if err := Register(nil); !errors.Is(err, sdk.ErrNilRegistrar) {
|
||||||
|
t.Fatalf("expected sdk.ErrNilRegistrar, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRegister_DuplicateErrors(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
reg := registry.New()
|
||||||
|
if err := Register(reg); err != nil {
|
||||||
|
t.Fatalf("first Register: %v", err)
|
||||||
|
}
|
||||||
|
if err := Register(reg); !errors.Is(err, registry.ErrDuplicate) {
|
||||||
|
t.Fatalf("expected registry.ErrDuplicate, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTools_ConcurrentDescribe(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
set, err := New(&config.Config{DataDir: t.TempDir()})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("New: %v", err)
|
||||||
|
}
|
||||||
|
// Exercise the read-only contract concurrently to catch any shared-state
|
||||||
|
// mutation (the dispatcher may call Parameters from many goroutines).
|
||||||
|
for _, tool := range set {
|
||||||
|
tool := tool
|
||||||
|
t.Run(tool.Name(), func(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
if _, err := tool.Parameters(); err != nil {
|
||||||
|
t.Fatalf("Parameters: %v", err)
|
||||||
|
}
|
||||||
|
if tool.Name() == "" || tool.Description() == "" {
|
||||||
|
t.Fatalf("empty identity for %v", tool)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,133 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"mime"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
// readLimited reads up to max bytes from r, returning the payload and whether it
|
||||||
|
// was truncated. A non-positive max falls back to the http response cap.
|
||||||
|
func readLimited(r io.Reader, max int64) ([]byte, bool, error) {
|
||||||
|
if max <= 0 {
|
||||||
|
max = defaultMaxHTTPBytes
|
||||||
|
}
|
||||||
|
data, err := io.ReadAll(io.LimitReader(r, max+1))
|
||||||
|
if err != nil {
|
||||||
|
return nil, false, err
|
||||||
|
}
|
||||||
|
truncated := int64(len(data)) > max
|
||||||
|
if truncated {
|
||||||
|
data = data[:max]
|
||||||
|
}
|
||||||
|
return data, truncated, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// validateContentType parses a caller-supplied Content-Type, returning the
|
||||||
|
// normalized media type. An empty value is allowed and yields "".
|
||||||
|
func validateContentType(ct string) (string, error) {
|
||||||
|
if strings.TrimSpace(ct) == "" {
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
mt, _, err := mime.ParseMediaType(ct)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("io: invalid content-type %q: %w", ct, err)
|
||||||
|
}
|
||||||
|
return mt, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// isTextContentType reports whether a media type is acceptable to fetch as text.
|
||||||
|
// An empty (unset) media type is accepted so a server that omits Content-Type is
|
||||||
|
// treated as text; binary media types are refused.
|
||||||
|
func isTextContentType(mt string) bool {
|
||||||
|
mt = strings.ToLower(strings.TrimSpace(mt))
|
||||||
|
if mt == "" {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
if strings.HasPrefix(mt, "text/") {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
switch mt {
|
||||||
|
case "application/json", "application/ld+json", "application/xml",
|
||||||
|
"application/xhtml+xml", "application/rss+xml", "application/atom+xml",
|
||||||
|
"application/x-www-form-urlencoded", "application/javascript",
|
||||||
|
"application/x-javascript", "image/svg+xml", "application/graphql":
|
||||||
|
return true
|
||||||
|
default:
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// isHTML reports whether a media type should be converted from HTML to text.
|
||||||
|
func isHTML(mt string) bool {
|
||||||
|
mt = strings.ToLower(strings.TrimSpace(mt))
|
||||||
|
return mt == "text/html" || mt == "application/xhtml+xml"
|
||||||
|
}
|
||||||
|
|
||||||
|
// htmlToText strips script/style/comment blocks and tags from an HTML document,
|
||||||
|
// decoding a few common entities and collapsing whitespace for readable text.
|
||||||
|
func htmlToText(src string) string {
|
||||||
|
s := dropBlocks(src)
|
||||||
|
var out strings.Builder
|
||||||
|
inTag := false
|
||||||
|
for i := 0; i < len(s); i++ {
|
||||||
|
switch c := s[i]; {
|
||||||
|
case c == '<':
|
||||||
|
inTag = true
|
||||||
|
case c == '>':
|
||||||
|
inTag = false
|
||||||
|
case !inTag:
|
||||||
|
out.WriteByte(c)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return normalizeSpace(htmlUnescape(out.String()))
|
||||||
|
}
|
||||||
|
|
||||||
|
// dropBlocks removes <script>, <style> and HTML comment regions.
|
||||||
|
func dropBlocks(s string) string {
|
||||||
|
for _, tag := range []string{"script", "style"} {
|
||||||
|
for {
|
||||||
|
start := strings.Index(s, "<"+tag)
|
||||||
|
if start < 0 {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
end := strings.Index(s, "</"+tag+">")
|
||||||
|
if end < 0 {
|
||||||
|
s = s[:start]
|
||||||
|
break
|
||||||
|
}
|
||||||
|
s = s[:start] + s[end+len(tag)+3:]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for {
|
||||||
|
start := strings.Index(s, "<!--")
|
||||||
|
if start < 0 {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
end := strings.Index(s, "-->")
|
||||||
|
if end < 0 {
|
||||||
|
s = s[:start]
|
||||||
|
break
|
||||||
|
}
|
||||||
|
s = s[:start] + s[end+3:]
|
||||||
|
}
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
|
||||||
|
// htmlUnescape replaces the small set of entities common in extracted text.
|
||||||
|
func htmlUnescape(s string) string {
|
||||||
|
return strings.NewReplacer(
|
||||||
|
"&", "&",
|
||||||
|
"<", "<",
|
||||||
|
">", ">",
|
||||||
|
""", "\"",
|
||||||
|
"'", "'",
|
||||||
|
" ", " ",
|
||||||
|
).Replace(s)
|
||||||
|
}
|
||||||
|
|
||||||
|
// normalizeSpace collapses runs of whitespace to single spaces and trims ends.
|
||||||
|
func normalizeSpace(s string) string {
|
||||||
|
return strings.Join(strings.Fields(s), " ")
|
||||||
|
}
|
||||||
@@ -0,0 +1,109 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
"mime"
|
||||||
|
"net/http"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Compile-time assertion that webFetchTool implements sdk.Tool.
|
||||||
|
var _ sdk.Tool = (*webFetchTool)(nil)
|
||||||
|
|
||||||
|
// webFetchTool fetches a URL over HTTP and returns its readable text, limited by
|
||||||
|
// size and content type, extracting text from HTML when appropriate.
|
||||||
|
type webFetchTool struct {
|
||||||
|
client *http.Client
|
||||||
|
allowedHosts []string
|
||||||
|
maxBytes int64
|
||||||
|
log *slog.Logger
|
||||||
|
}
|
||||||
|
|
||||||
|
func newWebFetchTool(t *tools) *webFetchTool {
|
||||||
|
return &webFetchTool{
|
||||||
|
client: t.client,
|
||||||
|
allowedHosts: t.allowedHosts,
|
||||||
|
maxBytes: t.maxFetchBytes,
|
||||||
|
log: t.log,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Name implements sdk.Tool.
|
||||||
|
func (w *webFetchTool) Name() string { return "web-fetch" }
|
||||||
|
|
||||||
|
// Description implements sdk.Tool.
|
||||||
|
func (w *webFetchTool) Description() string {
|
||||||
|
return "Fetch a URL and return its readable text, limited by size and content type."
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parameters implements sdk.Tool.
|
||||||
|
func (w *webFetchTool) Parameters() (map[string]any, error) {
|
||||||
|
return map[string]any{
|
||||||
|
"type": "object",
|
||||||
|
"properties": map[string]any{
|
||||||
|
"url": map[string]any{"type": "string", "description": "Absolute http(s) URL."},
|
||||||
|
"max_bytes": map[string]any{"type": "integer", "description": "Response size cap in bytes."},
|
||||||
|
},
|
||||||
|
"required": []any{"url"},
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Execute implements sdk.Tool.
|
||||||
|
func (w *webFetchTool) Execute(ctx context.Context, args any) (any, error) {
|
||||||
|
m, ok := args.(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("io: web-fetch: args must be an object")
|
||||||
|
}
|
||||||
|
rawURL := argString(m, "url")
|
||||||
|
if strings.TrimSpace(rawURL) == "" {
|
||||||
|
return nil, fmt.Errorf("io: web-fetch: url is required")
|
||||||
|
}
|
||||||
|
if _, err := validateRequestURL(ctx, rawURL, w.allowedHosts); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if w.client == nil {
|
||||||
|
return nil, fmt.Errorf("io: web-fetch: no underlying client")
|
||||||
|
}
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: web-fetch: build request: %w", err)
|
||||||
|
}
|
||||||
|
resp, err := w.client.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: web-fetch: do: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode != http.StatusOK {
|
||||||
|
return nil, fmt.Errorf("io: web-fetch: url returned status %d", resp.StatusCode)
|
||||||
|
}
|
||||||
|
mt := "text/plain"
|
||||||
|
if rawCT := resp.Header.Get("Content-Type"); rawCT != "" {
|
||||||
|
var err error
|
||||||
|
mt, _, err = mime.ParseMediaType(rawCT)
|
||||||
|
if err != nil || !isTextContentType(mt) {
|
||||||
|
return nil, fmt.Errorf("io: web-fetch: unexpected content type %q", rawCT)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
max := w.maxBytes
|
||||||
|
if v, ok := argInt(m, "max_bytes"); ok && v > 0 {
|
||||||
|
max = int64(v)
|
||||||
|
}
|
||||||
|
data, truncated, err := readLimited(resp.Body, max)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: web-fetch: read: %w", err)
|
||||||
|
}
|
||||||
|
text := string(data)
|
||||||
|
if isHTML(mt) {
|
||||||
|
text = htmlToText(text)
|
||||||
|
}
|
||||||
|
return map[string]any{
|
||||||
|
"url": rawURL,
|
||||||
|
"content_type": mt,
|
||||||
|
"text": text,
|
||||||
|
"truncated": truncated,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,108 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
// webFetchToolForTest builds a tools set with opts and returns the web-fetch
|
||||||
|
// tool.
|
||||||
|
func webFetchToolForTest(t *testing.T, opts ...Option) *webFetchTool {
|
||||||
|
t.Helper()
|
||||||
|
return findTool(t, buildTools(t, &config.Config{}, opts...), "web-fetch").(*webFetchTool)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWebFetch_TextExtraction(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
doc := `<!doctype html><html><head><title>T</title>` +
|
||||||
|
`<style>.x{color:red}</style><script>alert(1)</script></head>` +
|
||||||
|
`<body><h1>Hello</h1><p>World & more</p><!-- hidden --></body></html>`
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||||
|
_, _ = io.WriteString(w, doc)
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
wf := webFetchToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||||
|
|
||||||
|
got, err := wf.Execute(context.Background(), map[string]any{"url": srv.URL})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Execute: %v", err)
|
||||||
|
}
|
||||||
|
m := got.(map[string]any)
|
||||||
|
text := m["text"].(string)
|
||||||
|
if !strings.Contains(text, "Hello") || !strings.Contains(text, "World & more") {
|
||||||
|
t.Fatalf("text = %q", text)
|
||||||
|
}
|
||||||
|
if strings.Contains(text, "<") || strings.Contains(text, "alert") {
|
||||||
|
t.Fatalf("text still has tags/script: %q", text)
|
||||||
|
}
|
||||||
|
if m["content_type"] != "text/html" {
|
||||||
|
t.Fatalf("content_type = %q", m["content_type"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWebFetch_RejectsBinary(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "image/png")
|
||||||
|
_, _ = io.WriteString(w, "\x89PNG")
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
wf := webFetchToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||||
|
|
||||||
|
_, err := wf.Execute(context.Background(), map[string]any{"url": srv.URL})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "unexpected content type") {
|
||||||
|
t.Fatalf("expected content-type error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWebFetch_SSRFDenied(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
wf := webFetchToolForTest(t)
|
||||||
|
_, err := wf.Execute(context.Background(), map[string]any{"url": "http://localhost:9/"})
|
||||||
|
if !errors.Is(err, ErrSSRF) {
|
||||||
|
t.Fatalf("expected ErrSSRF, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWebFetch_SizeLimit(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/plain")
|
||||||
|
_, _ = io.WriteString(w, strings.Repeat("z", 200))
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
wf := webFetchToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||||
|
|
||||||
|
got, err := wf.Execute(context.Background(), map[string]any{"url": srv.URL, "max_bytes": 10})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Execute: %v", err)
|
||||||
|
}
|
||||||
|
m := got.(map[string]any)
|
||||||
|
if !m["truncated"].(bool) {
|
||||||
|
t.Fatal("expected truncated=true")
|
||||||
|
}
|
||||||
|
if len(m["text"].(string)) != 10 {
|
||||||
|
t.Fatalf("text len = %d", len(m["text"].(string)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWebFetch_Non200Status(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
http.Error(w, "bad", http.StatusNotFound)
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
wf := webFetchToolForTest(t, WithAllowedHosts(serverHost(srv)), WithHTTPClient(srv.Client()))
|
||||||
|
_, err := wf.Execute(context.Background(), map[string]any{"url": srv.URL})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "status 404") {
|
||||||
|
t.Fatalf("expected status-404 error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,161 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
"net/http"
|
||||||
|
"net/url"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Compile-time assertion that webSearchTool implements sdk.Tool.
|
||||||
|
var _ sdk.Tool = (*webSearchTool)(nil)
|
||||||
|
|
||||||
|
// webSearchTool queries a configured SearXNG-compatible (or custom) search
|
||||||
|
// endpoint and returns ranked results, capped in size and count.
|
||||||
|
type webSearchTool struct {
|
||||||
|
client *http.Client
|
||||||
|
endpoint string
|
||||||
|
allowedHosts []string
|
||||||
|
maxBytes int64
|
||||||
|
maxResults int
|
||||||
|
log *slog.Logger
|
||||||
|
}
|
||||||
|
|
||||||
|
func newWebSearchTool(t *tools) *webSearchTool {
|
||||||
|
return &webSearchTool{
|
||||||
|
client: t.client,
|
||||||
|
endpoint: t.searchURL,
|
||||||
|
allowedHosts: t.allowedHosts,
|
||||||
|
maxBytes: t.maxHTTPBytes,
|
||||||
|
maxResults: t.maxResults,
|
||||||
|
log: t.log,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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 via a configured SearXNG or custom search endpoint and return ranked results."
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parameters implements sdk.Tool.
|
||||||
|
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 := argString(m, "query")
|
||||||
|
if strings.TrimSpace(query) == "" {
|
||||||
|
return nil, fmt.Errorf("io: web-search: query is required")
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(s.endpoint) == "" {
|
||||||
|
return nil, fmt.Errorf("io: web-search: no search endpoint configured (set a SearXNG or custom URL)")
|
||||||
|
}
|
||||||
|
target, err := buildSearchURL(s.endpoint, query)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if _, err := validateRequestURL(ctx, target, s.allowedHosts); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if s.client == nil {
|
||||||
|
return nil, fmt.Errorf("io: web-search: no underlying client")
|
||||||
|
}
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: web-search: build request: %w", err)
|
||||||
|
}
|
||||||
|
resp, err := s.client.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: web-search: do: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode != http.StatusOK {
|
||||||
|
return nil, fmt.Errorf("io: web-search: endpoint returned status %d", resp.StatusCode)
|
||||||
|
}
|
||||||
|
data, truncated, err := readLimited(resp.Body, s.maxBytes)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("io: web-search: read: %w", err)
|
||||||
|
}
|
||||||
|
if truncated {
|
||||||
|
return nil, fmt.Errorf("io: web-search: response exceeds the %d-byte limit", s.maxBytes)
|
||||||
|
}
|
||||||
|
results, err := parseSearchResults(data)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
n := s.maxResults
|
||||||
|
if v, ok := argInt(m, "count"); ok && v > 0 {
|
||||||
|
n = v
|
||||||
|
}
|
||||||
|
if len(results) > n {
|
||||||
|
results = results[:n]
|
||||||
|
}
|
||||||
|
return map[string]any{
|
||||||
|
"query": query,
|
||||||
|
"count": len(results),
|
||||||
|
"results": results,
|
||||||
|
"truncated": truncated,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// buildSearchURL expands an endpoint against a query. It supports a "{query}"
|
||||||
|
// placeholder, otherwise appends q= and format=json for a SearXNG JSON API.
|
||||||
|
func buildSearchURL(endpoint, query string) (string, error) {
|
||||||
|
if strings.TrimSpace(endpoint) == "" {
|
||||||
|
return "", fmt.Errorf("io: web-search: no search endpoint configured")
|
||||||
|
}
|
||||||
|
q := url.QueryEscape(query)
|
||||||
|
if strings.Contains(endpoint, "{query}") {
|
||||||
|
return strings.ReplaceAll(endpoint, "{query}", q), nil
|
||||||
|
}
|
||||||
|
sep := "?"
|
||||||
|
if strings.Contains(endpoint, "?") {
|
||||||
|
sep = "&"
|
||||||
|
}
|
||||||
|
return endpoint + sep + "q=" + q + "&format=json", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// searchResult is one entry of a SearXNG-compatible results payload.
|
||||||
|
type searchResult struct {
|
||||||
|
Title string `json:"title"`
|
||||||
|
URL string `json:"url"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
Snippet string `json:"snippet"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// parseSearchResults decodes a {"results":[...]} payload, falling back to the
|
||||||
|
// snippet field when content is absent.
|
||||||
|
func parseSearchResults(data []byte) ([]searchResult, error) {
|
||||||
|
var payload struct {
|
||||||
|
Results []searchResult `json:"results"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(data, &payload); err != nil {
|
||||||
|
return nil, fmt.Errorf("io: web-search: decode results: %w", err)
|
||||||
|
}
|
||||||
|
for i := range payload.Results {
|
||||||
|
if payload.Results[i].Content == "" {
|
||||||
|
payload.Results[i].Content = payload.Results[i].Snippet
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return payload.Results, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,109 @@
|
|||||||
|
package io
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
// webSearchToolForTest builds a tools set with opts and returns the web-search
|
||||||
|
// tool.
|
||||||
|
func webSearchToolForTest(t *testing.T, opts ...Option) *webSearchTool {
|
||||||
|
t.Helper()
|
||||||
|
return findTool(t, buildTools(t, &config.Config{}, opts...), "web-search").(*webSearchTool)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWebSearch_NoEndpoint(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ws := webSearchToolForTest(t) // no endpoint configured
|
||||||
|
_, err := ws.Execute(context.Background(), map[string]any{"query": "cyrene"})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "no search endpoint") {
|
||||||
|
t.Fatalf("expected no-endpoint error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWebSearch_Query(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.URL.Query().Get("q") != "cyrene" {
|
||||||
|
t.Errorf("query param q = %q", r.URL.Query().Get("q"))
|
||||||
|
}
|
||||||
|
if r.URL.Query().Get("format") != "json" {
|
||||||
|
t.Errorf("format = %q", r.URL.Query().Get("format"))
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, _ = io.WriteString(w, `{"results":[
|
||||||
|
{"title":"A","url":"http://a.example","content":"alpha"},
|
||||||
|
{"title":"B","url":"http://b.example","snippet":"beta"}
|
||||||
|
]}`)
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
ws := webSearchToolForTest(t,
|
||||||
|
WithSearchEndpoint(srv.URL),
|
||||||
|
WithAllowedHosts(serverHost(srv)),
|
||||||
|
WithHTTPClient(srv.Client()),
|
||||||
|
)
|
||||||
|
|
||||||
|
got, err := ws.Execute(context.Background(), map[string]any{"query": "cyrene", "count": 2})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Execute: %v", err)
|
||||||
|
}
|
||||||
|
m := got.(map[string]any)
|
||||||
|
if m["count"].(int) != 2 {
|
||||||
|
t.Fatalf("count = %v", m["count"])
|
||||||
|
}
|
||||||
|
results := m["results"].([]searchResult)
|
||||||
|
if results[0].Title != "A" || results[0].Content != "alpha" {
|
||||||
|
t.Fatalf("result[0] = %+v", results[0])
|
||||||
|
}
|
||||||
|
if results[1].Content != "beta" {
|
||||||
|
t.Fatalf("result[1] snippet fallback = %+v", results[1])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWebSearch_SizeLimit(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
big := strings.Repeat("a", 64)
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
_, _ = io.WriteString(w, big)
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
ws := webSearchToolForTest(t,
|
||||||
|
WithSearchEndpoint(srv.URL),
|
||||||
|
WithAllowedHosts(serverHost(srv)),
|
||||||
|
WithHTTPClient(srv.Client()),
|
||||||
|
WithHTTPResponseLimit(16),
|
||||||
|
)
|
||||||
|
_, err := ws.Execute(context.Background(), map[string]any{"query": "cyrene"})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "exceeds the 16-byte limit") {
|
||||||
|
t.Fatalf("expected size-limit error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWebSearch_CountCapped(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, _ = io.WriteString(w, `{"results":[{"title":"A","url":"http://a","content":"x"}]}`)
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
ws := webSearchToolForTest(t,
|
||||||
|
WithSearchEndpoint(srv.URL),
|
||||||
|
WithAllowedHosts(serverHost(srv)),
|
||||||
|
WithHTTPClient(srv.Client()),
|
||||||
|
WithMaxResults(0), // ignored; default 10
|
||||||
|
WithMaxResults(1),
|
||||||
|
)
|
||||||
|
got, err := ws.Execute(context.Background(), map[string]any{"query": "cyrene"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Execute: %v", err)
|
||||||
|
}
|
||||||
|
if m := got.(map[string]any); m["count"].(int) != 1 {
|
||||||
|
t.Fatalf("count = %v", m["count"])
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in new issue
Block a user