feat(tools/io): 内置 IO/网络工具(file(路径安全)/http(SSRF防护)/web-search/web-fetch/knowledge)+ 大小截断

This commit is contained in:
Aska committed 2026-09-06 06:29:03 +08:00
1 parent 488cf9c79c
commit 3ce9489f8e
17 files changed
+2223

No files matched your search

+70
View File
@@ -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
}
}
+25
View File
@@ -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
+311
View File
@@ -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)
}
+187
View File
@@ -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")
}
}
+54
View File
@@ -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()
}
+127
View File
@@ -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
}
+189
View File
@@ -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)
}
}
+101
View File
@@ -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
}
+99
View File
@@ -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)
}
}
+101
View File
@@ -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
}
+232
View File
@@ -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
},
}
}
+107
View File
@@ -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)
}
}
})
}
}
+133
View File
@@ -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(
"&amp;", "&",
"&lt;", "<",
"&gt;", ">",
"&quot;", "\"",
"&#39;", "'",
"&nbsp;", " ",
).Replace(s)
}
// normalizeSpace collapses runs of whitespace to single spaces and trims ends.
func normalizeSpace(s string) string {
return strings.Join(strings.Fields(s), " ")
}
+109
View File
@@ -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
}
+108
View File
@@ -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 &amp; 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)
}
}
+161
View File
@@ -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
}
+109
View File
@@ -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"])
}
}