Cyrene_v2/internal/capability/registry/policy_store_test.go

package registry

import (
	"context"
	"database/sql"
	"errors"
	"io"
	"log/slog"
	"os"
	"path/filepath"
	"testing"

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

// 本文件锁定裁决 1.3/10.4.2（工具策略落盘）：迁移 0012_tool_policy 的 up/down
// 可逆、以及「重启后策略仍在」——后者用「关连接、重建注册表、再 Load」模拟
// 一次真实重启。

// errStoreDown 是内存桩注入的落盘失败。
var errStoreDown = errors.New("policy store down")

// testStoreLogger 返回一个静默 logger：存储层测试不关心日志输出。
func testStoreLogger() *slog.Logger {
	return slog.New(slog.NewTextHandler(io.Discard, nil))
}

// openMigrated 开一个（文件或 :memory:）库并跑完全部迁移。
// 用文件而非 :memory: 的用例是因为「重启仍在」要跨连接读同一份数据。
func openMigrated(t *testing.T, path string) *sql.DB {
	t.Helper()
	ctx := context.Background()
	conn, err := db.Open(ctx, path, testStoreLogger())
	if err != nil {
		t.Fatalf("db.Open(%s): %v", path, err)
	}
	if err := db.MigrateUp(ctx, conn, testStoreLogger()); err != nil {
		conn.Close()
		t.Fatalf("MigrateUp: %v", err)
	}
	t.Cleanup(func() { conn.Close() })
	return conn
}

// toolPoliciesTableExists 报告 tool_policies 表是否存在。
func toolPoliciesTableExists(t *testing.T, ctx context.Context, conn *sql.DB) bool {
	t.Helper()
	var name string
	err := conn.QueryRowContext(ctx,
		`SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'tool_policies'`).Scan(&name)
	if errors.Is(err, sql.ErrNoRows) {
		return false
	}
	if err != nil {
		t.Fatalf("query sqlite_master: %v", err)
	}
	return true
}

// copyFile 复制一个文件（副本迁移验证用）。
func copyFile(t *testing.T, src, dst string) {
	t.Helper()
	data, err := os.ReadFile(src)
	if err != nil {
		t.Fatalf("read %s: %v", src, err)
	}
	if err := os.WriteFile(dst, data, 0o600); err != nil {
		t.Fatalf("write %s: %v", dst, err)
	}
}

// TestToolPolicyMigration_UpDownUpOnDataDirCopy 在 data/cyrene.db 的**副本**上
// 跑迁移 0012 的 up → down → up，证明可逆且幂等（绝不碰本体）。
// 副本不存在时跳过（例如干净检出的 CI），此时 :memory: 的用例仍然覆盖同一迁移。
func TestToolPolicyMigration_UpDownUpOnDataDirCopy(t *testing.T) {
	ctx := context.Background()
	src := filepath.Join("..", "..", "..", "data", "cyrene.db")
	if _, err := os.Stat(src); err != nil {
		t.Skipf("data/cyrene.db 不可读，跳过副本迁移验证: %v", err)
	}
	dst := filepath.Join(t.TempDir(), "cyrene-copy.db")
	copyFile(t, src, dst)
	// WAL 模式下已提交内容可能还在 -wal 里，一并带过去，副本才是完整快照。
	for _, suffix := range []string{"-wal", "-shm"} {
		if _, err := os.Stat(src + suffix); err == nil {
			copyFile(t, src+suffix, dst+suffix)
		}
	}

	conn, err := db.Open(ctx, dst, testStoreLogger())
	if err != nil {
		t.Fatalf("db.Open(copy): %v", err)
	}
	defer conn.Close()

	if err := db.MigrateUp(ctx, conn, testStoreLogger()); err != nil {
		t.Fatalf("MigrateUp(copy): %v", err)
	}
	if !toolPoliciesTableExists(t, ctx, conn) {
		t.Fatal("tool_policies 应在 MigrateUp 后存在")
	}
	if _, err := conn.ExecContext(ctx,
		`INSERT INTO tool_policies (name, allow_in_group) VALUES ('copied_tool', 1)`); err != nil {
		t.Fatalf("insert into copy: %v", err)
	}

	// down：只回滚最后一步（= 0012）。
	if err := db.MigrateDown(ctx, conn, 1, testStoreLogger()); err != nil {
		t.Fatalf("MigrateDown(copy, 1): %v", err)
	}
	if toolPoliciesTableExists(t, ctx, conn) {
		t.Fatal("tool_policies 应在 MigrateDown 后消失")
	}

	// 再 up：可逆且幂等。
	if err := db.MigrateUp(ctx, conn, testStoreLogger()); err != nil {
		t.Fatalf("MigrateUp(copy, 2nd): %v", err)
	}
	if !toolPoliciesTableExists(t, ctx, conn) {
		t.Fatal("tool_policies 应在第二次 MigrateUp 后重新出现")
	}
	var rows int
	if err := conn.QueryRowContext(ctx, `SELECT COUNT(*) FROM tool_policies`).Scan(&rows); err != nil {
		t.Fatalf("count rows: %v", err)
	}
	if rows != 0 {
		t.Fatalf("回滚再升级后应有 0 行，得到 %d 行（回滚未清空？）", rows)
	}
}

// TestToolPolicyMigration_UpDownOnMemory 在干净库上锁定同一迁移的 up/down
// （不依赖 data/ 目录是否存在）。
func TestToolPolicyMigration_UpDownOnMemory(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	conn := openMigrated(t, ":memory:")

	if !toolPoliciesTableExists(t, ctx, conn) {
		t.Fatal("tool_policies 应在 MigrateUp 后存在")
	}
	if err := db.MigrateDown(ctx, conn, 1, testStoreLogger()); err != nil {
		t.Fatalf("MigrateDown: %v", err)
	}
	if toolPoliciesTableExists(t, ctx, conn) {
		t.Fatal("tool_policies 应在 MigrateDown 后消失")
	}
}

// TestRegistry_PolicyStoreSurvivesRestart 锁定「重启后策略仍在」：
// 第一次运行 Save，第二次运行（新连接 + 新注册表 + 新 store）Load 回来。
func TestRegistry_PolicyStoreSurvivesRestart(t *testing.T) {
	ctx := context.Background()
	path := filepath.Join(t.TempDir(), "cyrene.db")

	first := openMigrated(t, path)
	store, err := NewSQLitePolicyStore(first, testStoreLogger())
	if err != nil {
		t.Fatalf("NewSQLitePolicyStore: %v", err)
	}
	reg := New()
	if err := reg.AttachPolicyStore(ctx, store); err != nil {
		t.Fatalf("AttachPolicyStore: %v", err)
	}
	for _, name := range []string{"group_tool", "other"} {
		if err := reg.RegisterTool(regTool{name: name}); err != nil {
			t.Fatalf("RegisterTool(%s): %v", name, err)
		}
	}
	// 初始值必须是默认（两个布尔 false = owner-only）。
	if pol, err := reg.Policy("group_tool"); err != nil || pol.AllowInGroup || pol.AllowInStranger {
		t.Fatalf("initial policy = %+v (err %v), want both false", pol, err)
	}
	if err := reg.SetPolicy("group_tool", ToolPolicy{AllowInGroup: true}); err != nil {
		t.Fatalf("SetPolicy: %v", err)
	}
	if err := first.Close(); err != nil { // 模拟进程退出
		t.Fatalf("close first connection: %v", err)
	}

	second := openMigrated(t, path)
	store2, err := NewSQLitePolicyStore(second, testStoreLogger())
	if err != nil {
		t.Fatalf("NewSQLitePolicyStore(2): %v", err)
	}
	reg2 := New()
	if err := reg2.AttachPolicyStore(ctx, store2); err != nil {
		t.Fatalf("AttachPolicyStore(2): %v", err)
	}
	// 重启后注册表里还没有能力，只有 Load 出的策略快照……
	if _, err := reg2.Policy("group_tool"); !errors.Is(err, ErrNotFound) {
		t.Fatalf("Policy before re-register error = %v, want ErrNotFound", err)
	}
	// ……能力注册进来的那一刻套上落盘值（MCP 热装等晚注册路径同理）。
	if err := reg2.RegisterTool(regTool{name: "group_tool"}); err != nil {
		t.Fatalf("RegisterTool(2): %v", err)
	}
	pol, err := reg2.Policy("group_tool")
	if err != nil {
		t.Fatalf("Policy(2): %v", err)
	}
	if !pol.AllowInGroup || pol.AllowInStranger {
		t.Fatalf("policy after restart = %+v, want {AllowInGroup:true AllowInStranger:false}", pol)
	}
	// 没配过的能力仍是默认（tool_policies 里没有它的行）。
	if err := reg2.RegisterTool(regTool{name: "other"}); err != nil {
		t.Fatalf("RegisterTool(other): %v", err)
	}
	if pol, err := reg2.Policy("other"); err != nil || pol.AllowInGroup || pol.AllowInStranger {
		t.Fatalf("unconfigured policy = %+v (err %v), want both false", pol, err)
	}
}

// TestRegistry_AttachPolicyStoreAppliesToRegisteredCaps 锁定反向顺序：
// 能力先注册、策略后 Load（AttachPolicyStore 在 wireRuntime 才启用的场景），
// 已注册能力也必须立刻拿到落盘值。
func TestRegistry_AttachPolicyStoreAppliesToRegisteredCaps(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	seed := newMemoryPolicyStore()
	seed.marks["stranger_tool"] = ToolPolicy{AllowInStranger: true}

	reg := New()
	if err := reg.RegisterTool(regTool{name: "stranger_tool"}); err != nil {
		t.Fatalf("RegisterTool: %v", err)
	}
	if err := reg.AttachPolicyStore(ctx, seed); err != nil {
		t.Fatalf("AttachPolicyStore: %v", err)
	}
	pol, err := reg.Policy("stranger_tool")
	if err != nil {
		t.Fatalf("Policy: %v", err)
	}
	if pol.AllowInGroup || !pol.AllowInStranger {
		t.Fatalf("policy after attach = %+v, want {AllowInGroup:false AllowInStranger:true}", pol)
	}
}

// TestRegistry_SetPolicyPersistFailureKeepsMemory 锁定「先落库、后改内存」：
// 落盘失败时内存保持旧值且返回错误，绝不出现「面板说成功、重启却复位」。
func TestRegistry_SetPolicyPersistFailureKeepsMemory(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	store := newMemoryPolicyStore()
	reg := New()
	if err := reg.RegisterTool(regTool{name: "calc"}); err != nil {
		t.Fatalf("RegisterTool: %v", err)
	}
	if err := reg.AttachPolicyStore(ctx, store); err != nil {
		t.Fatalf("AttachPolicyStore: %v", err)
	}
	store.failSave = true
	if err := reg.SetPolicy("calc", ToolPolicy{AllowInGroup: true}); err == nil {
		t.Fatal("SetPolicy with failing store = nil error, want error")
	}
	pol, err := reg.Policy("calc")
	if err != nil {
		t.Fatalf("Policy: %v", err)
	}
	if pol.AllowInGroup || pol.AllowInStranger {
		t.Fatalf("policy after failed persist = %+v, want unchanged (both false)", pol)
	}
	// 单布尔入口（SetAllowInGroup/SetAllowInStranger）走同一条落库路径。
	if err := reg.SetAllowInGroup("calc", true); err == nil {
		t.Fatal("SetAllowInGroup with failing store = nil error, want error")
	}
}

// memoryPolicyStore 是 PolicyStore 的内存桩：Load/Save 语义与 SQLite 实现一致
// （只保存写过的能力），failSave 用于注入落盘失败。
type memoryPolicyStore struct {
	marks    map[string]ToolPolicy
	failSave bool
}

func newMemoryPolicyStore() *memoryPolicyStore {
	return &memoryPolicyStore{marks: make(map[string]ToolPolicy)}
}

func (m *memoryPolicyStore) Load(_ context.Context) (map[string]ToolPolicy, error) {
	out := make(map[string]ToolPolicy, len(m.marks))
	for k, v := range m.marks {
		out[k] = v
	}
	return out, nil
}

func (m *memoryPolicyStore) Save(_ context.Context, name string, p ToolPolicy) error {
	if m.failSave {
		return errStoreDown
	}
	m.marks[name] = p
	return nil
}