Cyrene_v2/internal/capability/mcp/schemacache.go

package mcp

import (
	"strings"
	"sync"

	"github.com/google/jsonschema-go/jsonschema"
)

// SchemaCache 缓存 tools/list 里已解析的 input schema（design/09 §1.1）。
//
// 键 = 服务器名 + 工具名；收到 notifications/tools/list_changed 时按服务器整表失效。
// 命中时 CallTool 直接用缓存 schema 校验参数并补默认值，不重复拉取/解析 schema。
// 并发安全；命中/未命中计数对调用方可读，便于验证缓存真的生效。
type SchemaCache struct {
	mu      sync.Mutex
	entries map[string]*cachedSchema
	hits    int
	misses  int
}

// cachedSchema 是一条缓存：原始 schema（模型可见的形状）+ 已解析结果。
// resolved 为 nil 表示 schema 解析失败，调用时降级为不校验。
type cachedSchema struct {
	raw      map[string]any
	resolved *jsonschema.Resolved
}

// NewSchemaCache 构建空缓存。
func NewSchemaCache() *SchemaCache {
	return &SchemaCache{entries: map[string]*cachedSchema{}}
}

// Put 解析并缓存一条 input schema。解析失败时仍缓存原始 schema（校验降级）并返回
// 解析错误，由调用方决定是否告警。
func (c *SchemaCache) Put(server, tool string, schema map[string]any) error {
	if schema == nil {
		schema = emptyObjectSchema()
	}
	entry := &cachedSchema{raw: schema}
	resolved, err := resolveSchema(schema)
	entry.resolved = resolved

	c.mu.Lock()
	c.entries[cacheKey(server, tool)] = entry
	c.mu.Unlock()
	return err
}

// Get 返回缓存的原始 input schema。
func (c *SchemaCache) Get(server, tool string) (map[string]any, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	entry, ok := c.entries[cacheKey(server, tool)]
	if !ok {
		return nil, false
	}
	return entry.raw, true
}

// Prepare 返回「已补默认值、已校验」的参数：
//
//   - 命中缓存：用缓存 schema 补默认值 + 校验（计一次命中），不再重复拉取/解析；
//   - 未命中：有 fallback（本次 tools/list 发现结果）则解析后缓存再用；没有可用
//     schema 时原样返回（计一次未命中），即首次工具清单到达前的直接调用降级为不校验。
//
// 传给服务器的是返回的参数值；args 为 nil 时按空对象处理。
func (c *SchemaCache) Prepare(server, tool string, fallback, args map[string]any) (map[string]any, error) {
	entry := c.lookup(server, tool)
	if entry == nil {
		if fallback == nil {
			return args, nil
		}
		// 解析失败不阻断调用：Put 已把原始 schema 放进缓存，resolved 为 nil 即不校验。
		_ = c.Put(server, tool, fallback)
		entry = c.lookup(server, tool)
	}
	if entry == nil || entry.resolved == nil {
		return args, nil
	}
	if args == nil {
		args = map[string]any{}
	}
	if err := applySchema(entry.resolved, args); err != nil {
		return nil, err
	}
	return args, nil
}

// lookup 取一条缓存并记账（命中/未命中）。
func (c *SchemaCache) lookup(server, tool string) *cachedSchema {
	c.mu.Lock()
	defer c.mu.Unlock()
	entry, ok := c.entries[cacheKey(server, tool)]
	if !ok {
		c.misses++
		return nil
	}
	c.hits++
	return entry
}

// Invalidate 让某台服务器的全部缓存失效；工具清单变化时调用。
func (c *SchemaCache) Invalidate(server string) {
	prefix := server + keySeparator
	c.mu.Lock()
	defer c.mu.Unlock()
	for k := range c.entries {
		if strings.HasPrefix(k, prefix) {
			delete(c.entries, k)
		}
	}
}

// Clear 清空所有缓存。
func (c *SchemaCache) Clear() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.entries = map[string]*cachedSchema{}
}

// Len 返回当前缓存条目数。
func (c *SchemaCache) Len() int {
	c.mu.Lock()
	defer c.mu.Unlock()
	return len(c.entries)
}

// Stats 返回命中次数与未命中次数。
func (c *SchemaCache) Stats() (hits, misses int) {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.hits, c.misses
}

// keySeparator 分隔服务器名与工具名，避免不同服务器同名工具互相覆盖。
const keySeparator = "\x00"

// cacheKey 拼接缓存键。
func cacheKey(server, tool string) string {
	return server + keySeparator + tool
}