Agent-Native/runtime/tools/registry.go

package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	"git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/runtime/checkpoint"
	"git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/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 {
	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)
		}
	}
	if len(spec.Properties) > 0 {
		for key := range obj {
			if _, ok := spec.Properties[key]; !ok {
				return fmt.Errorf("unknown argument %s", key)
			}
		}
	}
	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 {
	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