feat(tools): add Tool definitions, Registry pattern, and filesystem sandbox with path validation
This commit is contained in:
parent
9f33d6df93
commit
c25243a690
5 changed files with 522 additions and 0 deletions
65
pkg/tools/registry.go
Normal file
65
pkg/tools/registry.go
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"sync"
|
||||
)
|
||||
|
||||
// registry is the default implementation of Registry.
|
||||
type registry struct {
|
||||
mu sync.RWMutex
|
||||
tools map[string]Tool
|
||||
order []string
|
||||
}
|
||||
|
||||
// NewRegistry returns a new empty registry.
|
||||
func NewRegistry() Registry {
|
||||
return ®istry{
|
||||
tools: make(map[string]Tool),
|
||||
}
|
||||
}
|
||||
|
||||
func (r *registry) Register(t Tool) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
if _, exists := r.tools[t.Name]; exists {
|
||||
return ErrDuplicateTool
|
||||
}
|
||||
|
||||
r.tools[t.Name] = t
|
||||
r.order = append(r.order, t.Name)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *registry) Get(name string) (Tool, bool) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
t, ok := r.tools[name]
|
||||
return t, ok
|
||||
}
|
||||
|
||||
func (r *registry) List() []Tool {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
out := make([]Tool, len(r.order))
|
||||
for i, name := range r.order {
|
||||
out[i] = r.tools[name]
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (r *registry) Filter(predicate func(Tool) bool) []Tool {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
var result []Tool
|
||||
for _, name := range r.order {
|
||||
t := r.tools[name]
|
||||
if predicate(t) {
|
||||
result = append(result, t)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
114
pkg/tools/registry_test.go
Normal file
114
pkg/tools/registry_test.go
Normal file
|
|
@ -0,0 +1,114 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNewRegistry(t *testing.T) {
|
||||
reg := NewRegistry()
|
||||
if reg == nil {
|
||||
t.Fatal("expected non-nil registry")
|
||||
}
|
||||
if len(reg.List()) != 0 {
|
||||
t.Errorf("expected empty registry, got %d tools", len(reg.List()))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterAndGet(t *testing.T) {
|
||||
reg := NewRegistry()
|
||||
tool := Tool{Name: "test", Description: "test tool"}
|
||||
if err := reg.Register(tool); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
got, ok := reg.Get("test")
|
||||
if !ok {
|
||||
t.Fatal("expected tool to be found")
|
||||
}
|
||||
if got.Name != "test" {
|
||||
t.Errorf("expected name 'test', got %q", got.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterDuplicate(t *testing.T) {
|
||||
reg := NewRegistry()
|
||||
tool := Tool{Name: "test"}
|
||||
if err := reg.Register(tool); err != nil {
|
||||
t.Fatalf("unexpected error on first register: %v", err)
|
||||
}
|
||||
if err := reg.Register(tool); err == nil {
|
||||
t.Error("expected error for duplicate tool")
|
||||
}
|
||||
}
|
||||
|
||||
func TestList(t *testing.T) {
|
||||
reg := NewRegistry()
|
||||
tools := []Tool{
|
||||
{Name: "a"},
|
||||
{Name: "b"},
|
||||
{Name: "c"},
|
||||
}
|
||||
for _, tw := range tools {
|
||||
if err := reg.Register(tw); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
list := reg.List()
|
||||
if len(list) != 3 {
|
||||
t.Errorf("expected 3 tools, got %d", len(list))
|
||||
}
|
||||
}
|
||||
|
||||
func TestFilter(t *testing.T) {
|
||||
reg := NewRegistry()
|
||||
tools := []Tool{
|
||||
{Name: "read", Permission: Allow},
|
||||
{Name: "write", Permission: Ask},
|
||||
{Name: "delete", Permission: Deny},
|
||||
}
|
||||
for _, tw := range tools {
|
||||
if err := reg.Register(tw); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
allowed := reg.Filter(func(t Tool) bool {
|
||||
return t.Permission == Allow
|
||||
})
|
||||
if len(allowed) != 1 {
|
||||
t.Errorf("expected 1 allowed tool, got %d", len(allowed))
|
||||
}
|
||||
if len(allowed) > 0 && allowed[0].Name != "read" {
|
||||
t.Errorf("expected 'read' tool, got %q", allowed[0].Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetNotFound(t *testing.T) {
|
||||
reg := NewRegistry()
|
||||
_, ok := reg.Get("nonexistent")
|
||||
if ok {
|
||||
t.Error("expected false for nonexistent tool")
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrentRegister(t *testing.T) {
|
||||
reg := NewRegistry()
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 100; i++ {
|
||||
wg.Add(1)
|
||||
go func(n int) {
|
||||
defer wg.Done()
|
||||
name := fmt.Sprintf("tool_%d", n)
|
||||
_ = reg.Register(Tool{Name: name})
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
list := reg.List()
|
||||
if len(list) != 100 {
|
||||
t.Errorf("expected 100 tools, got %d", len(list))
|
||||
}
|
||||
}
|
||||
115
pkg/tools/sandbox/sandbox.go
Normal file
115
pkg/tools/sandbox/sandbox.go
Normal file
|
|
@ -0,0 +1,115 @@
|
|||
package sandbox
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/tools"
|
||||
)
|
||||
|
||||
// ErrSandboxViolation is returned when a tool call violates sandbox rules.
|
||||
var ErrSandboxViolation = fmt.Errorf("sandbox violation")
|
||||
|
||||
// Sandbox wraps os.Root to enforce filesystem boundaries for tool calls.
|
||||
type Sandbox struct {
|
||||
root *os.Root
|
||||
path string
|
||||
}
|
||||
|
||||
// NewSandbox creates a new Sandbox rooted at the given directory.
|
||||
// The directory is created if it does not exist.
|
||||
func NewSandbox(rootDir string) (*Sandbox, error) {
|
||||
if err := os.MkdirAll(rootDir, 0755); err != nil {
|
||||
return nil, fmt.Errorf("creating sandbox root: %w", err)
|
||||
}
|
||||
|
||||
absRoot, err := filepath.Abs(rootDir)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("resolving absolute path: %w", err)
|
||||
}
|
||||
|
||||
root, err := os.OpenRoot(absRoot)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("opening sandbox root: %w", err)
|
||||
}
|
||||
|
||||
return &Sandbox{root: root, path: absRoot}, nil
|
||||
}
|
||||
|
||||
// ValidateToolCall checks if a tool call's arguments reference paths outside the sandbox.
|
||||
// It returns nil if the call is allowed, or an error explaining the violation.
|
||||
func (s *Sandbox) ValidateToolCall(tool tools.Tool, call llm.ToolCall) error {
|
||||
args := make(map[string]interface{})
|
||||
if len(call.Arguments) > 0 {
|
||||
if err := json.Unmarshal(call.Arguments, &args); err != nil {
|
||||
return fmt.Errorf("parsing arguments: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
for _, path := range extractPaths(args) {
|
||||
if err := s.validatePath(path); err != nil {
|
||||
return fmt.Errorf("%w: %s", ErrSandboxViolation, err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// validatePath checks that the given path resolves inside the sandbox root.
|
||||
func (s *Sandbox) validatePath(path string) error {
|
||||
// Absolute paths are rejected (they escape the sandbox by definition)
|
||||
if filepath.IsAbs(path) {
|
||||
return fmt.Errorf("absolute paths not allowed: %s", path)
|
||||
}
|
||||
|
||||
// Resolve to absolute path relative to sandbox root
|
||||
abs := filepath.Join(s.path, path)
|
||||
|
||||
// Clean the path to normalize it
|
||||
abs = filepath.Clean(abs)
|
||||
|
||||
// Check if the path is inside the root
|
||||
if !strings.HasPrefix(abs, s.path+string(filepath.Separator)) && abs != s.path {
|
||||
return fmt.Errorf("path escapes sandbox: %s (root: %s)", abs, s.path)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// extractPaths collects all path-like values from the arguments map.
|
||||
func extractPaths(args map[string]interface{}) []string {
|
||||
var paths []string
|
||||
for _, v := range args {
|
||||
switch val := v.(type) {
|
||||
case string:
|
||||
if isPathLike(val) {
|
||||
paths = append(paths, val)
|
||||
}
|
||||
case []interface{}:
|
||||
for _, item := range val {
|
||||
if str, ok := item.(string); ok && isPathLike(str) {
|
||||
paths = append(paths, str)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return paths
|
||||
}
|
||||
|
||||
// isPathLike checks if a string looks like a filesystem path.
|
||||
func isPathLike(s string) bool {
|
||||
// Must start with / or ./ or ../ or contain a file extension
|
||||
return strings.HasPrefix(s, "/") ||
|
||||
strings.HasPrefix(s, "./") ||
|
||||
strings.HasPrefix(s, "../") ||
|
||||
strings.Contains(s, ".") && strings.Contains(s, "/")
|
||||
}
|
||||
|
||||
// Open returns the underlying os.Root for testing.
|
||||
func (s *Sandbox) Open() *os.Root {
|
||||
return s.root
|
||||
}
|
||||
149
pkg/tools/sandbox/sandbox_test.go
Normal file
149
pkg/tools/sandbox/sandbox_test.go
Normal file
|
|
@ -0,0 +1,149 @@
|
|||
package sandbox_test
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/tools"
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/tools/sandbox"
|
||||
)
|
||||
|
||||
func TestNewSandbox(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
|
||||
sb, err := sandbox.NewSandbox(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if sb == nil {
|
||||
t.Fatal("expected non-nil sandbox")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidatePath_Allowed(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
sb, err := sandbox.NewSandbox(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// Create a file inside the sandbox
|
||||
testFile := filepath.Join(dir, "test.txt")
|
||||
if err := os.WriteFile(testFile, []byte("hello"), 0644); err != nil {
|
||||
t.Fatalf("creating test file: %v", err)
|
||||
}
|
||||
|
||||
call := llm.ToolCall{
|
||||
Name: "read_file",
|
||||
Arguments: json.RawMessage(`{"path": "test.txt"}`),
|
||||
}
|
||||
|
||||
err = sb.ValidateToolCall(tools.Tool{}, call)
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidatePath_EscapesSandbox(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
sb, err := sandbox.NewSandbox(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
call := llm.ToolCall{
|
||||
Name: "read_file",
|
||||
Arguments: json.RawMessage(`{"path": "/etc/passwd"}`),
|
||||
}
|
||||
|
||||
err = sb.ValidateToolCall(tools.Tool{}, call)
|
||||
if err == nil {
|
||||
t.Fatal("expected sandbox violation error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidatePath_SymlinkEscape(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
|
||||
// Create a symlink that escapes the sandbox
|
||||
escapeDir := t.TempDir()
|
||||
symlinkPath := filepath.Join(dir, "link")
|
||||
if err := os.Symlink(escapeDir, symlinkPath); err != nil {
|
||||
t.Fatalf("creating symlink: %v", err)
|
||||
}
|
||||
|
||||
sb, err := sandbox.NewSandbox(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// The sandbox should allow symlinks inside the root
|
||||
call := llm.ToolCall{
|
||||
Name: "read_file",
|
||||
Arguments: json.RawMessage(`{"path": "link"}`),
|
||||
}
|
||||
|
||||
err = sb.ValidateToolCall(tools.Tool{}, call)
|
||||
if err != nil {
|
||||
t.Fatalf("expected symlink to be allowed (path is inside sandbox), got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidatePath_NonExistentFile(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
sb, err := sandbox.NewSandbox(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// Non-existent file inside sandbox should be allowed
|
||||
call := llm.ToolCall{
|
||||
Name: "read_file",
|
||||
Arguments: json.RawMessage(`{"path": "nonexistent.txt"}`),
|
||||
}
|
||||
|
||||
err = sb.ValidateToolCall(tools.Tool{}, call)
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error for non-existent file, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidatePath_InvalidJSON(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
sb, err := sandbox.NewSandbox(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
call := llm.ToolCall{
|
||||
Name: "read_file",
|
||||
Arguments: json.RawMessage(`not json`),
|
||||
}
|
||||
|
||||
err = sb.ValidateToolCall(tools.Tool{}, call)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for invalid JSON")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidatePath_NoPathsInArgs(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
sb, err := sandbox.NewSandbox(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// Arguments with no paths should not cause errors
|
||||
call := llm.ToolCall{
|
||||
Name: "math_add",
|
||||
Arguments: json.RawMessage(`{"a": 5, "b": 10}`),
|
||||
}
|
||||
|
||||
err = sb.ValidateToolCall(tools.Tool{}, call)
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error, got: %v", err)
|
||||
}
|
||||
}
|
||||
79
pkg/tools/types.go
Normal file
79
pkg/tools/types.go
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// Permission controls whether a tool can be executed without user approval.
|
||||
type Permission int
|
||||
|
||||
const (
|
||||
Allow Permission = iota
|
||||
Ask
|
||||
Deny
|
||||
)
|
||||
|
||||
func (p Permission) String() string {
|
||||
switch p {
|
||||
case Allow:
|
||||
return "allow"
|
||||
case Ask:
|
||||
return "ask"
|
||||
case Deny:
|
||||
return "deny"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// Tool is a function callable by the LLM with JSON Schema input.
|
||||
type Tool struct {
|
||||
Name string
|
||||
Description string
|
||||
InputSchema json.RawMessage // JSON Schema draft-07+
|
||||
Required []string
|
||||
Handler ToolHandler
|
||||
Permission Permission
|
||||
Examples []ToolExample
|
||||
}
|
||||
|
||||
// ToolHandler is the function type that implements a tool.
|
||||
// The handler receives raw JSON arguments and returns a result.
|
||||
type ToolHandler func(ctx context.Context, args json.RawMessage) (ToolResult, error)
|
||||
|
||||
// ToolResult is returned by a ToolHandler.
|
||||
type ToolResult struct {
|
||||
Content string
|
||||
IsError bool
|
||||
Metadata map[string]string
|
||||
Artifacts []Artifact
|
||||
}
|
||||
|
||||
// Artifact represents a file or data artifact produced by a tool.
|
||||
type Artifact struct {
|
||||
Path string
|
||||
Content []byte
|
||||
MIME string
|
||||
}
|
||||
|
||||
// ToolExample provides few-shot examples for the LLM to improve tool usage.
|
||||
type ToolExample struct {
|
||||
Input map[string]interface{}
|
||||
Output string
|
||||
}
|
||||
|
||||
// Registry manages tool registration and lookup.
|
||||
type Registry interface {
|
||||
Register(tool Tool) error
|
||||
Get(name string) (Tool, bool)
|
||||
List() []Tool
|
||||
Filter(predicate func(Tool) bool) []Tool
|
||||
}
|
||||
|
||||
// ErrToolNotFound is returned when a tool is not found in the registry.
|
||||
var ErrToolNotFound = fmt.Errorf("tool not found")
|
||||
|
||||
// ErrDuplicateTool is returned when trying to register a tool with an existing name.
|
||||
var ErrDuplicateTool = fmt.Errorf("duplicate tool name")
|
||||
Loading…
Reference in a new issue