368 lines
9.3 KiB
Go
368 lines
9.3 KiB
Go
package tools
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"agent-native/runtime/argschema"
|
|
"agent-native/runtime/checkpoint"
|
|
"agent-native/runtime/filestate"
|
|
)
|
|
|
|
// 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)
|
|
}
|
|
|
|
const (
|
|
// ReplaySafe 表示恢复时可以再执行一次。ReplayNever 是缺省，不重跑。
|
|
ReplaySafe = "safe"
|
|
ReplayNever = "never"
|
|
)
|
|
|
|
// Replayer 是可选声明。未实现或返回其他值都按 never。
|
|
type Replayer interface {
|
|
Replay() string
|
|
}
|
|
|
|
// ReplayOf 读取工具的重放声明。
|
|
func ReplayOf(t Tool) string {
|
|
r, ok := t.(Replayer)
|
|
if !ok || r.Replay() != ReplaySafe {
|
|
return ReplayNever
|
|
}
|
|
return ReplaySafe
|
|
}
|
|
|
|
// 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
|
|
|
|
const (
|
|
// ExposureDirect 进入模型看到的工具清单。ExposureDeferred 只通过 tool_search 与 call_tool。
|
|
ExposureDirect = "direct"
|
|
ExposureDeferred = "deferred"
|
|
)
|
|
|
|
// Exposer 声明工具是否直接出现在清单里。未实现视为直接。
|
|
type Exposer interface {
|
|
Exposure() string
|
|
}
|
|
|
|
// ExposureOf 读取暴露级别。
|
|
func ExposureOf(t Tool) string {
|
|
if e, ok := t.(Exposer); ok && e.Exposure() == ExposureDeferred {
|
|
return ExposureDeferred
|
|
}
|
|
return ExposureDirect
|
|
}
|
|
|
|
// 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
|
|
// hookMu 串行 Before、After 与 Audit，避免并行工具调用交错写审计或重入钩子。
|
|
hookMu sync.Mutex
|
|
mu sync.Mutex
|
|
// restricted 为真时，后注册的工具也必须命中 allowExact / allowPrefix。
|
|
restricted bool
|
|
allowExact map[string]bool
|
|
allowPrefix []string
|
|
activated map[string]bool
|
|
activatePath string
|
|
}
|
|
|
|
func NewRegistry(root string) *Registry {
|
|
return &Registry{root: root, tools: map[string]Tool{}}
|
|
}
|
|
|
|
func (r *Registry) Root() string { return r.root }
|
|
|
|
func (r *Registry) Has(name string) bool {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
_, ok := r.tools[name]
|
|
return ok
|
|
}
|
|
|
|
// Replay 返回已注册工具的重放声明。未知工具是 never。
|
|
func (r *Registry) Replay(name string) string {
|
|
if r == nil {
|
|
return ReplayNever
|
|
}
|
|
r.mu.Lock()
|
|
t, ok := r.tools[name]
|
|
r.mu.Unlock()
|
|
if !ok {
|
|
return ReplayNever
|
|
}
|
|
return ReplayOf(t)
|
|
}
|
|
|
|
func (r *Registry) Register(t Tool) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
if r.restricted && !r.keepLocked(t.Name()) {
|
|
return
|
|
}
|
|
if _, ok := r.tools[t.Name()]; !ok {
|
|
r.order = append(r.order, t.Name())
|
|
}
|
|
r.tools[t.Name()] = t
|
|
if ExposureOf(t) == ExposureDeferred {
|
|
r.ensureMetaLocked()
|
|
}
|
|
}
|
|
|
|
func (r *Registry) Decls() []Decl {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
out := make([]Decl, 0, len(r.order))
|
|
for _, name := range r.order {
|
|
t := r.tools[name]
|
|
if ExposureOf(t) == ExposureDeferred {
|
|
continue
|
|
}
|
|
out = append(out, Decl{Name: t.Name(), Description: t.Description(), Parameters: t.Parameters()})
|
|
}
|
|
return out
|
|
}
|
|
|
|
// Restrict 把注册表收窄到 allow 列出的工具（工具策略的唯一执行点之一：
|
|
// 调用方声明 allowlist，注册表执行收窄）。条目支持两种形式：
|
|
//
|
|
// 精确名（read_file）或命名空间通配（mcp__github__*，匹配该前缀全部工具）。
|
|
//
|
|
// 任一条目未命中任何已注册工具时返回错误（fail-closed，不静默忽略配置错误）。
|
|
// 空列表不收窄；可重复调用（在已收窄集合上继续求交）。
|
|
func (r *Registry) Restrict(allow []string) error {
|
|
if len(allow) == 0 {
|
|
return nil
|
|
}
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
exact := make(map[string]bool, len(allow))
|
|
prefixes := make([]string, 0)
|
|
for _, entry := range allow {
|
|
if name, ok := strings.CutSuffix(entry, "*"); ok {
|
|
prefixes = append(prefixes, name)
|
|
continue
|
|
}
|
|
exact[entry] = true
|
|
}
|
|
var unknown []string
|
|
for _, name := range allow {
|
|
if _, ok := strings.CutSuffix(name, "*"); ok {
|
|
continue
|
|
}
|
|
if _, ok := r.tools[name]; !ok {
|
|
unknown = append(unknown, name)
|
|
}
|
|
}
|
|
for _, p := range prefixes {
|
|
if !r.hasPrefix(p) {
|
|
unknown = append(unknown, p+"*")
|
|
}
|
|
}
|
|
if len(unknown) > 0 {
|
|
return fmt.Errorf("tool allowlist names not registered: %s", strings.Join(unknown, ", "))
|
|
}
|
|
r.restricted = true
|
|
r.allowExact = exact
|
|
r.allowPrefix = prefixes
|
|
var order []string
|
|
for _, name := range r.order {
|
|
if r.keepLocked(name) {
|
|
order = append(order, name)
|
|
} else {
|
|
delete(r.tools, name)
|
|
}
|
|
}
|
|
r.order = order
|
|
return nil
|
|
}
|
|
|
|
func (r *Registry) keepLocked(name string) bool {
|
|
if name == "tool_search" || name == "call_tool" {
|
|
return true
|
|
}
|
|
if r.allowExact[name] {
|
|
return true
|
|
}
|
|
for _, p := range r.allowPrefix {
|
|
if strings.HasPrefix(name, p) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (r *Registry) hasPrefix(prefix string) bool {
|
|
for _, name := range r.order {
|
|
if strings.HasPrefix(name, prefix) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// 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 {
|
|
r.mu.Lock()
|
|
t, ok := r.tools[name]
|
|
r.mu.Unlock()
|
|
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 {
|
|
r.hookMu.Lock()
|
|
block, reason := r.Before(ctx, name, args)
|
|
r.hookMu.Unlock()
|
|
if 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 {
|
|
r.hookMu.Lock()
|
|
res = r.After(ctx, name, args, res)
|
|
r.hookMu.Unlock()
|
|
}
|
|
r.audit(name, false, res.IsError, "")
|
|
return res
|
|
}
|
|
|
|
func (r *Registry) audit(name string, blocked, isError bool, note string) {
|
|
r.hookMu.Lock()
|
|
defer r.hookMu.Unlock()
|
|
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 {
|
|
return argschema.Validate(schema, args)
|
|
}
|
|
|
|
// WorkspaceTools registers the document workspace tool set.
|
|
func WorkspaceTools(root string) *Registry {
|
|
return WorkspaceToolsTracked(root, nil)
|
|
}
|
|
|
|
// WorkspaceToolsTracked 与 WorkspaceTools 相同，但改写已存在的文件前要先在本轮读过。
|
|
func WorkspaceToolsTracked(root string, files *filestate.Tracker) *Registry {
|
|
return WorkspaceToolsCheckpointed(root, files, nil, nil)
|
|
}
|
|
|
|
// CheckpointedTools 是会写检查点的工具名集合——回滚资格的唯一声明处
|
|
// （回滚只撤销这三类写入；shell/python/其他工具永不在集）。集合由
|
|
// WorkspaceToolsCheckpointed 的构造使用，并有测试钉住，防静默漂移。
|
|
var CheckpointedTools = []string{"write_file", "edit_file", "multi_edit"}
|
|
|
|
// Checkpointed 报告一个工具名是否有回滚资格（供宿主/测试查询）。
|
|
func Checkpointed(name string) bool {
|
|
for _, n := range CheckpointedTools {
|
|
if n == name {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// WorkspaceToolsCheckpointed 在写、改文件前把改前内容交给检查点。log 为空则不记录。
|
|
func WorkspaceToolsCheckpointed(root string, files *filestate.Tracker, log *checkpoint.Log, anchor func() string) *Registry {
|
|
mark := writeMark(log, root, anchor)
|
|
r := NewRegistry(root)
|
|
r.Register(&readFile{root: root, files: files})
|
|
r.Register(&writeFile{root: root, files: files, mark: mark})
|
|
r.Register(&editFile{root: root, files: files, mark: mark})
|
|
r.Register(&multiEditTool{root: root, files: files, mark: mark})
|
|
r.Register(&listDir{root: root})
|
|
r.Register(&grepTool{root: root})
|
|
r.Register(&globTool{root: root})
|
|
return r
|
|
}
|
|
|
|
func writeMark(log *checkpoint.Log, root string, anchor func() string) func(rel string, before []byte, existed bool, after []byte) {
|
|
if log == nil {
|
|
return nil
|
|
}
|
|
return func(rel string, before []byte, existed bool, after []byte) {
|
|
id := ""
|
|
if anchor != nil {
|
|
id = anchor()
|
|
}
|
|
_ = log.Record(id, root, rel, before, existed, after)
|
|
}
|
|
}
|
|
|
|
func rejectStaleWrite(files *filestate.Tracker, path string) string {
|
|
if files == nil {
|
|
return ""
|
|
}
|
|
if err := files.CheckWrite(path); err != nil {
|
|
return err.Error()
|
|
}
|
|
return ""
|
|
}
|
|
|
|
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
|