Agent-Native/runtime/tools/registry.go

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) Has(name string) bool {
	_, ok := r.tools[name]
	return ok
}

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
}

// Restrict 把注册表收窄到 allow 列出的工具（工具策略的唯一执行点之一：
// 调用方声明 allowlist，注册表执行收窄）。条目支持两种形式：
//
//	精确名（read_file）或命名空间通配（mcp__github__*，匹配该前缀全部工具）。
//
// 任一条目未命中任何已注册工具时返回错误（fail-closed，不静默忽略配置错误）。
// 空列表不收窄；可重复调用（在已收窄集合上继续求交）。
func (r *Registry) Restrict(allow []string) error {
	if len(allow) == 0 {
		return nil
	}
	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, ", "))
	}
	keep := func(name string) bool {
		if exact[name] {
			return true
		}
		for _, p := range prefixes {
			if strings.HasPrefix(name, p) {
				return true
			}
		}
		return false
	}
	var order []string
	for _, name := range r.order {
		if keep(name) {
			order = append(order, name)
		} else {
			delete(r.tools, name)
		}
	}
	r.order = order
	return nil
}

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 {
	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(&multiEditTool{root: root})
	r.Register(&listDir{root: root})
	r.Register(&grepTool{root: root})
	r.Register(&globTool{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