Cyrene_v2/internal/capability/registry/policy_store.go

package registry

import (
	"context"
	"database/sql"
	"fmt"
	"log/slog"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

// PolicyStore 持久化工具级可达性标记（design/11 §3、design/09 §2.1；裁决 1.3/10.4.2）。
//
// 载体是 SQLite 的 tool_policies 表（迁移 0012_tool_policy）；*SQLitePolicyStore
// 是生产实现，测试可换成内存桩。实现必须对并发使用安全。
type PolicyStore interface {
	// Load 返回全部已落盘的策略，键为能力名。从未配置过的能力不出现在结果里
	// ——调用方按默认值（两个布尔均 false = owner-only）处理，即 fail-closed。
	Load(ctx context.Context) (map[string]ToolPolicy, error)
	// Save 写入一个能力的策略：不存在则插入，存在则更新（并把 updated_at
	// 推进到当前时间）。name 为空是非法调用，返回错误。
	Save(ctx context.Context, name string, p ToolPolicy) error
}

// 策略持久化的 SQL（全参数化，项目规范 §十）。表由迁移 0012_tool_policy 建立；
// upsert 以唯一索引 idx_tool_policies_name 为冲突目标，因此同一能力永远只有一行。
const (
	selectToolPoliciesSQL = `SELECT name, allow_in_group, allow_in_stranger FROM tool_policies`
	upsertToolPolicySQL   = `
INSERT INTO tool_policies (name, allow_in_group, allow_in_stranger)
VALUES (?, ?, ?)
ON CONFLICT(name) DO UPDATE SET
    allow_in_group    = excluded.allow_in_group,
    allow_in_stranger = excluded.allow_in_stranger,
    updated_at        = strftime('%Y-%m-%dT%H:%M:%fZ','now')`
)

// AttachPolicyStore 把持久化载体接到注册表（design/11 §3；裁决 1.3/10.4.2）：
//
//  1. 立即 Load 全部落盘策略，套用到**已注册**的能力上（重启后策略仍在）；
//  2. 记住这份快照与 store，之后 Register 进来的能力（MCP 热装、插件刷新）
//     会用落盘值覆盖其初始标记；
//  3. 此后 SetPolicy/SetAllowInGroup/SetAllowInStranger 先落库、成功后才改内存。
//
// store 为 nil 是合法的空操作（注册表保持纯内存，即既有行为）。Load 失败时
// 返回错误且不动任何内存状态——调用方可以选择忽略（默认值就是最严的
// owner-only，fail-closed），但绝不能把「没有读出策略」当成「策略全开」。
func (r *Registry) AttachPolicyStore(ctx context.Context, store PolicyStore) error {
	if r == nil || store == nil {
		return nil
	}
	loaded, err := store.Load(ctx)
	if err != nil {
		return fmt.Errorf("registry: attach policy store: %w", err)
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.store = store
	r.persisted = loaded
	for name, pol := range loaded {
		cap, ok := r.caps[name]
		if !ok {
			continue // 能力尚未注册：Register 时会用 r.persisted 套上
		}
		cap.Policy = pol
		r.caps[name] = cap
	}
	return nil
}

// persist 把一条策略写进已装上的 store（未装时空操作）。它刻意在注册表锁之外
// 执行 SQL，调用方只在它成功后才改内存——即「先落库、后改内存」。
//
// 它用 context.Background()：策略写入口是管理端点经 CapabilityStore 窄接口
// （gateway/contracts.go）调用的既有签名 SetPolicy(name, p)，没有 ctx 可传；
// 本地 SQLite 单行 upsert 是短操作，不值得为此改契约。
func (r *Registry) persist(name string, p ToolPolicy) error {
	store := r.storeSnapshot()
	if store == nil {
		return nil
	}
	if err := store.Save(context.Background(), name, p); err != nil {
		return fmt.Errorf("registry: persist policy %q: %w", name, err)
	}
	return nil
}

// storeSnapshot 在锁内读一次 store 指针，避免与 AttachPolicyStore 的写入竞态。
func (r *Registry) storeSnapshot() PolicyStore {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return r.store
}

// remember 记录刚落盘成功的策略（调用方必须已持有写锁）。它使「卸下再重注册」
// 的能力（插件刷新）仍拿到最新标记，而不是注册时的零值。
func (r *Registry) remember(name string, p ToolPolicy) {
	if r.store == nil {
		return
	}
	if r.persisted == nil {
		r.persisted = make(map[string]ToolPolicy)
	}
	r.persisted[name] = p
}

// SQLitePolicyStore 把工具策略持久化在 SQLite 的 tool_policies 表上。
// 它凭借底层的 *sql.DB 对并发使用安全。
type SQLitePolicyStore struct {
	db  *sql.DB
	log *slog.Logger
}

// NewSQLitePolicyStore 在 db 上构建 SQLite 策略存储。db 为 nil 时返回错误。
//
// 它刻意不自建 DDL：迁移 0012_tool_policy 是这份 schema 的唯一权威路径，
// 因此表缺失时 Load/Save 会返回明确的错误（调用方据此保持最严默认值），
// 而不是悄悄降级成内存态、让管理面板的「保存成功」变成重启即复位的假成功。
func NewSQLitePolicyStore(db *sql.DB, log *slog.Logger) (*SQLitePolicyStore, error) {
	if db == nil {
		return nil, fmt.Errorf("registry: policy store: nil database")
	}
	return &SQLitePolicyStore{db: db, log: logger.WithSource(log, "capability/registry")}, nil
}

// Load 返回 tool_policies 的全部行。空表返回空 map（不是错误）：那就是
// 「从未配置过任何标记」的诚实状态，所有能力按默认 owner-only 处理。
func (s *SQLitePolicyStore) Load(ctx context.Context) (map[string]ToolPolicy, error) {
	if s == nil || s.db == nil {
		return nil, fmt.Errorf("registry: policy store: not configured")
	}
	rows, err := s.db.QueryContext(ctx, selectToolPoliciesSQL)
	if err != nil {
		return nil, fmt.Errorf("registry: load tool policies: %w", err)
	}
	defer rows.Close()

	out := make(map[string]ToolPolicy)
	for rows.Next() {
		var (
			name            string
			allowInGroup    bool
			allowInStranger bool
		)
		if err := rows.Scan(&name, &allowInGroup, &allowInStranger); err != nil {
			return nil, fmt.Errorf("registry: scan tool policy: %w", err)
		}
		out[name] = ToolPolicy{AllowInGroup: allowInGroup, AllowInStranger: allowInStranger}
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("registry: iterate tool policies: %w", err)
	}
	return out, nil
}

// Save upsert name 的策略。两条布尔标记本身不是敏感信息（能力名 + 两个布尔），
// 因此错误信息只带能力名；日志也不打印任何参数外的内容。
func (s *SQLitePolicyStore) Save(ctx context.Context, name string, p ToolPolicy) error {
	if s == nil || s.db == nil {
		return fmt.Errorf("registry: policy store: not configured")
	}
	if name == "" {
		return fmt.Errorf("registry: save tool policy: %w", ErrInvalidArgument)
	}
	if _, err := s.db.ExecContext(ctx, upsertToolPolicySQL, name, p.AllowInGroup, p.AllowInStranger); err != nil {
		return fmt.Errorf("registry: save tool policy %q: %w", name, err)
	}
	return nil
}