From 022b914c6652d0b9349e2d90512a05b214bf15c5 Mon Sep 17 00:00:00 2001 From: weichen Date: Mon, 21 Sep 2026 14:49:13 +0000 Subject: [PATCH] feat(agentx): jail workspace tools and audit hooks Keep read, write, edit, list, and grep inside the workspace root, including symlink targets, and let hooks block or rewrite a call. --- internal/agentx/tools/builtin.go | 198 +++++++++++++++++++++++++++ internal/agentx/tools/jail.go | 85 ++++++++++++ internal/agentx/tools/registry.go | 202 ++++++++++++++++++++++++++++ internal/agentx/tools/tools_test.go | 105 +++++++++++++++ 4 files changed, 590 insertions(+) create mode 100644 internal/agentx/tools/builtin.go create mode 100644 internal/agentx/tools/jail.go create mode 100644 internal/agentx/tools/registry.go create mode 100644 internal/agentx/tools/tools_test.go diff --git a/internal/agentx/tools/builtin.go b/internal/agentx/tools/builtin.go new file mode 100644 index 00000000..7fef7ec2 --- /dev/null +++ b/internal/agentx/tools/builtin.go @@ -0,0 +1,198 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" +) + +type readFile struct{ root string } + +func (t *readFile) Name() string { return "read_file" } +func (t *readFile) Description() string { + return "Read a UTF-8 text file inside the workspace." +} +func (t *readFile) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","required":["path"],"properties":{"path":{"type":"string"}}}`) +} +func (t *readFile) Execute(_ context.Context, args json.RawMessage) (Result, error) { + rel, err := readArg(args, "path") + if err != nil { + return Result{}, err + } + path, err := mustPath(t.root, rel) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + b, err := os.ReadFile(path) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + if len(b) > maxReadBytes { + b = b[:maxReadBytes] + } + return Result{Content: string(b)}, nil +} + +type writeFile struct{ root string } + +func (t *writeFile) Name() string { return "write_file" } +func (t *writeFile) Description() string { + return "Create or replace a UTF-8 text file inside the workspace." +} +func (t *writeFile) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","required":["path","content"],"properties":{"path":{"type":"string"},"content":{"type":"string"}}}`) +} +func (t *writeFile) Execute(_ context.Context, args json.RawMessage) (Result, error) { + var obj struct { + Path string `json:"path"` + Content string `json:"content"` + } + if err := json.Unmarshal(args, &obj); err != nil { + return Result{}, err + } + path, err := mustPath(t.root, obj.Path) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + if err := os.WriteFile(path, []byte(obj.Content), 0o644); err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + return Result{Content: "wrote " + obj.Path}, nil +} + +type editFile struct{ root string } + +func (t *editFile) Name() string { return "edit_file" } +func (t *editFile) Description() string { + return "Replace one unique exact text span in a workspace file." +} +func (t *editFile) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","required":["path","old","new"],"properties":{"path":{"type":"string"},"old":{"type":"string"},"new":{"type":"string"}}}`) +} +func (t *editFile) Execute(_ context.Context, args json.RawMessage) (Result, error) { + var obj struct { + Path string `json:"path"` + Old string `json:"old"` + New string `json:"new"` + } + if err := json.Unmarshal(args, &obj); err != nil { + return Result{}, err + } + if obj.Old == "" { + return Result{Content: "old text is empty", IsError: true}, nil + } + path, err := mustPath(t.root, obj.Path) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + b, err := os.ReadFile(path) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + text := string(b) + n := strings.Count(text, obj.Old) + if n != 1 { + return Result{Content: fmt.Sprintf("expected one match, found %d", n), IsError: true}, nil + } + updated := strings.Replace(text, obj.Old, obj.New, 1) + if err := os.WriteFile(path, []byte(updated), 0o644); err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + return Result{Content: "edited " + obj.Path}, nil +} + +type listDir struct{ root string } + +func (t *listDir) Name() string { return "list_dir" } +func (t *listDir) Description() string { + return "List names in a workspace directory." +} +func (t *listDir) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}}}`) +} +func (t *listDir) Execute(_ context.Context, args json.RawMessage) (Result, error) { + rel, _ := readArg(args, "path") + if rel == "" { + rel = "." + } + path, err := mustPath(t.root, rel) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + entries, err := os.ReadDir(path) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + names := make([]string, 0, len(entries)) + for _, e := range entries { + name := e.Name() + if e.IsDir() { + name += "/" + } + names = append(names, name) + } + return Result{Content: strings.Join(names, "\n")}, nil +} + +type grepTool struct{ root string } + +func (t *grepTool) Name() string { return "grep" } +func (t *grepTool) Description() string { + return "Search workspace text files for a literal substring." +} +func (t *grepTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","required":["pattern"],"properties":{"pattern":{"type":"string"},"path":{"type":"string"}}}`) +} +func (t *grepTool) Execute(_ context.Context, args json.RawMessage) (Result, error) { + var obj struct { + Pattern string `json:"pattern"` + Path string `json:"path"` + } + if err := json.Unmarshal(args, &obj); err != nil { + return Result{}, err + } + if obj.Path == "" { + obj.Path = "." + } + start, err := mustPath(t.root, obj.Path) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + var hits []string + err = filepath.WalkDir(start, func(path string, d os.DirEntry, walkErr error) error { + if walkErr != nil || d.IsDir() { + return nil + } + if len(hits) >= 50 { + return fmt.Errorf("stop") + } + b, err := os.ReadFile(path) + if err != nil || len(b) > maxReadBytes { + return nil + } + rel, _ := filepath.Rel(t.root, path) + for i, line := range strings.Split(string(b), "\n") { + if strings.Contains(line, obj.Pattern) { + hits = append(hits, fmt.Sprintf("%s:%d:%s", filepath.ToSlash(rel), i+1, line)) + if len(hits) >= 50 { + return fmt.Errorf("stop") + } + } + } + return nil + }) + if err != nil && err.Error() != "stop" { + return Result{Content: err.Error(), IsError: true}, nil + } + if len(hits) == 0 { + return Result{Content: "no matches"}, nil + } + return Result{Content: strings.Join(hits, "\n")}, nil +} diff --git a/internal/agentx/tools/jail.go b/internal/agentx/tools/jail.go new file mode 100644 index 00000000..f7162e4e --- /dev/null +++ b/internal/agentx/tools/jail.go @@ -0,0 +1,85 @@ +package tools + +import ( + "fmt" + "os" + "path/filepath" + "strings" +) + +// Resolve joins rel to root and rejects paths that escape the workspace, +// including symlink targets outside root. Missing files are allowed when their +// existing ancestors stay inside root. +func Resolve(root, rel string) (string, error) { + rootAbs, err := filepath.Abs(root) + if err != nil { + return "", err + } + rootAbs = filepath.Clean(rootAbs) + resolvedRoot, err := filepath.EvalSymlinks(rootAbs) + if err != nil { + return "", err + } + rel = strings.TrimSpace(rel) + if rel == "" { + rel = "." + } + if filepath.IsAbs(rel) { + inside, err := filepath.Rel(rootAbs, filepath.Clean(rel)) + if err != nil || escapes(inside) { + return "", fmt.Errorf("path escapes workspace") + } + rel = inside + } + rel = filepath.Clean(rel) + if escapes(rel) { + return "", fmt.Errorf("path escapes workspace") + } + if rel == "." { + return resolvedRoot, nil + } + current := resolvedRoot + parts := strings.Split(rel, string(os.PathSeparator)) + for i, part := range parts { + if part == "" || part == "." { + continue + } + if part == ".." { + return "", fmt.Errorf("path escapes workspace") + } + next := filepath.Join(current, part) + info, err := os.Lstat(next) + if os.IsNotExist(err) { + tail := filepath.Join(parts[i:]...) + return filepath.Join(current, tail), nil + } + if err != nil { + return "", err + } + if info.Mode()&os.ModeSymlink != 0 { + target, err := filepath.EvalSymlinks(next) + if err != nil { + return "", fmt.Errorf("path escapes workspace") + } + relTarget, err := filepath.Rel(resolvedRoot, target) + if err != nil || escapes(relTarget) { + return "", fmt.Errorf("path escapes workspace") + } + current = target + continue + } + current = next + } + relFinal, err := filepath.Rel(resolvedRoot, current) + if err != nil || escapes(relFinal) { + return "", fmt.Errorf("path escapes workspace") + } + return current, nil +} + +func escapes(rel string) bool { + if rel == ".." { + return true + } + return strings.HasPrefix(rel, ".."+string(os.PathSeparator)) +} diff --git a/internal/agentx/tools/registry.go b/internal/agentx/tools/registry.go new file mode 100644 index 00000000..07974775 --- /dev/null +++ b/internal/agentx/tools/registry.go @@ -0,0 +1,202 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "io" + "strings" + "time" +) + +// Result is the tool payload returned to the model. +type Result struct { + Content string + IsError bool +} + +// Tool is one executable function with a JSON schema. +type Tool interface { + Name() string + Description() string + Parameters() json.RawMessage + Execute(ctx context.Context, args json.RawMessage) (Result, error) +} + +// BeforeHook can block a call after schema validation. +type BeforeHook func(ctx context.Context, name string, args json.RawMessage) (block bool, reason string) + +// AfterHook can rewrite a tool result before it is returned. +type AfterHook func(ctx context.Context, name string, args json.RawMessage, result Result) Result + +// Registry executes jailed tools and appends an audit line per call. +type Registry struct { + root string + tools map[string]Tool + order []string + Before BeforeHook + After AfterHook + Audit io.Writer +} + +func NewRegistry(root string) *Registry { + return &Registry{root: root, tools: map[string]Tool{}} +} + +func (r *Registry) Root() string { return r.root } + +func (r *Registry) Register(t Tool) { + if _, ok := r.tools[t.Name()]; !ok { + r.order = append(r.order, t.Name()) + } + r.tools[t.Name()] = t +} + +func (r *Registry) Decls() []Decl { + out := make([]Decl, 0, len(r.order)) + for _, name := range r.order { + t := r.tools[name] + out = append(out, Decl{Name: t.Name(), Description: t.Description(), Parameters: t.Parameters()}) + } + return out +} + +// Decl is the schema view of a tool. +type Decl struct { + Name string + Description string + Parameters json.RawMessage +} + +// Execute validates arguments, runs hooks, and audits the outcome. +func (r *Registry) Execute(ctx context.Context, name string, args json.RawMessage) Result { + t, ok := r.tools[name] + if !ok { + res := Result{Content: "unknown tool " + name, IsError: true} + r.audit(name, false, true, "unknown") + return res + } + if len(args) == 0 { + args = []byte(`{}`) + } + if err := validateArgs(t.Parameters(), args); err != nil { + res := Result{Content: err.Error(), IsError: true} + r.audit(name, false, true, "schema") + return res + } + if r.Before != nil { + if block, reason := r.Before(ctx, name, args); block { + if reason == "" { + reason = "blocked" + } + res := Result{Content: reason, IsError: true} + r.audit(name, true, true, reason) + return res + } + } + res, err := t.Execute(ctx, args) + if err != nil { + res = Result{Content: err.Error(), IsError: true} + } + if r.After != nil { + res = r.After(ctx, name, args, res) + } + r.audit(name, false, res.IsError, "") + return res +} + +func (r *Registry) audit(name string, blocked, isError bool, note string) { + if r.Audit == nil { + return + } + fmt.Fprintf(r.Audit, "%s\t%s\tblocked=%t\terror=%t\t%s\n", time.Now().UTC().Format(time.RFC3339Nano), name, blocked, isError, note) +} + +func validateArgs(schema, args json.RawMessage) error { + var spec struct { + Type string `json:"type"` + Required []string `json:"required"` + Properties map[string]json.RawMessage `json:"properties"` + } + if len(schema) > 0 { + if err := json.Unmarshal(schema, &spec); err != nil { + return fmt.Errorf("invalid tool schema") + } + } + var obj map[string]any + if err := json.Unmarshal(args, &obj); err != nil { + return fmt.Errorf("arguments must be a JSON object") + } + if spec.Type != "" && spec.Type != "object" { + return fmt.Errorf("unsupported schema type %s", spec.Type) + } + for _, req := range spec.Required { + v, ok := obj[req] + if !ok || v == nil { + return fmt.Errorf("missing required argument %s", req) + } + if s, ok := v.(string); ok && strings.TrimSpace(s) == "" { + return fmt.Errorf("missing required argument %s", req) + } + } + for key, raw := range spec.Properties { + v, ok := obj[key] + if !ok { + continue + } + var prop struct { + Type string `json:"type"` + } + if err := json.Unmarshal(raw, &prop); err != nil || prop.Type == "" { + continue + } + if !matchesType(prop.Type, v) { + return fmt.Errorf("argument %s must be %s", key, prop.Type) + } + } + return nil +} + +func matchesType(kind string, v any) bool { + switch kind { + case "string": + _, ok := v.(string) + return ok + case "number": + _, ok := v.(float64) + return ok + case "integer": + f, ok := v.(float64) + return ok && f == float64(int64(f)) + case "boolean": + _, ok := v.(bool) + return ok + default: + return true + } +} + +// WorkspaceTools registers the document workspace tool set. +func WorkspaceTools(root string) *Registry { + r := NewRegistry(root) + r.Register(&readFile{root: root}) + r.Register(&writeFile{root: root}) + r.Register(&editFile{root: root}) + r.Register(&listDir{root: root}) + r.Register(&grepTool{root: root}) + return r +} + +func mustPath(root, rel string) (string, error) { + return Resolve(root, rel) +} + +func readArg(args json.RawMessage, key string) (string, error) { + var obj map[string]string + if err := json.Unmarshal(args, &obj); err != nil { + return "", err + } + return obj[key], nil +} + +const maxReadBytes = 200 << 10 diff --git a/internal/agentx/tools/tools_test.go b/internal/agentx/tools/tools_test.go new file mode 100644 index 00000000..6f1f0c9d --- /dev/null +++ b/internal/agentx/tools/tools_test.go @@ -0,0 +1,105 @@ +package tools + +import ( + "bytes" + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestSchemaRejectsMissingAndWrongType(t *testing.T) { + dir := t.TempDir() + reg := WorkspaceTools(dir) + missing := reg.Execute(context.Background(), "read_file", []byte(`{}`)) + if !missing.IsError || !strings.Contains(missing.Content, "path") { + t.Fatalf("missing %#v", missing) + } + wrong := reg.Execute(context.Background(), "read_file", []byte(`{"path":1}`)) + if !wrong.IsError || !strings.Contains(wrong.Content, "string") { + t.Fatalf("type %#v", wrong) + } +} + +func TestPathEscapeAndSymlink(t *testing.T) { + dir := t.TempDir() + outside := t.TempDir() + secret := filepath.Join(outside, "secret.txt") + if err := os.WriteFile(secret, []byte("hidden"), 0o644); err != nil { + t.Fatal(err) + } + link := filepath.Join(dir, "leak") + if err := os.Symlink(secret, link); err != nil { + t.Fatal(err) + } + reg := WorkspaceTools(dir) + up := reg.Execute(context.Background(), "read_file", []byte(`{"path":"../secret.txt"}`)) + if !up.IsError { + t.Fatalf("parent escape succeeded: %s", up.Content) + } + via := reg.Execute(context.Background(), "read_file", []byte(`{"path":"leak"}`)) + if !via.IsError || strings.Contains(via.Content, "hidden") { + t.Fatalf("symlink escape: %#v", via) + } + absOutside := reg.Execute(context.Background(), "write_file", []byte(`{"path":"`+secret+`","content":"no"}`)) + if !absOutside.IsError { + t.Fatal("absolute outside write succeeded") + } +} + +func TestWriteReadEditGrep(t *testing.T) { + dir := t.TempDir() + reg := WorkspaceTools(dir) + wrote := reg.Execute(context.Background(), "write_file", []byte(`{"path":"docs/a.md","content":"alpha beta"}`)) + if wrote.IsError { + t.Fatal(wrote.Content) + } + got := reg.Execute(context.Background(), "read_file", []byte(`{"path":"docs/a.md"}`)) + if got.IsError || got.Content != "alpha beta" { + t.Fatalf("read %#v", got) + } + edited := reg.Execute(context.Background(), "edit_file", []byte(`{"path":"docs/a.md","old":"beta","new":"gamma"}`)) + if edited.IsError { + t.Fatal(edited.Content) + } + found := reg.Execute(context.Background(), "grep", []byte(`{"pattern":"gamma"}`)) + if found.IsError || !strings.Contains(found.Content, "docs/a.md:1:") { + t.Fatalf("grep %#v", found) + } + listed := reg.Execute(context.Background(), "list_dir", []byte(`{"path":"docs"}`)) + if !strings.Contains(listed.Content, "a.md") { + t.Fatalf("list %#v", listed) + } +} + +func TestBeforeHookBlocksAndAudits(t *testing.T) { + dir := t.TempDir() + reg := WorkspaceTools(dir) + var audit bytes.Buffer + reg.Audit = &audit + reg.Before = func(ctx context.Context, name string, args json.RawMessage) (bool, string) { + return true, "policy" + } + res := reg.Execute(context.Background(), "list_dir", []byte(`{}`)) + if !res.IsError || res.Content != "policy" { + t.Fatalf("hook %#v", res) + } + if !strings.Contains(audit.String(), "blocked=true") { + t.Fatalf("audit %s", audit.String()) + } +} + +func TestAfterHookRewrites(t *testing.T) { + dir := t.TempDir() + reg := WorkspaceTools(dir) + reg.After = func(ctx context.Context, name string, args json.RawMessage, result Result) Result { + result.Content = "rewritten" + return result + } + res := reg.Execute(context.Background(), "list_dir", []byte(`{}`)) + if res.Content != "rewritten" { + t.Fatalf("%#v", res) + } +}