Cyrene_v2/plugins/tools/io/file_claims_rig_test.go

package io

import (
	"context"
	"errors"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/tool"
	"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
)

// 本文件是 file_claims_test.go（真实插件 + 真实分发器的互斥回归）的夹具：记录版占用表、
// 带闸门的工具包装、以及等待与断言助手。夹具只搭台子，不写断言语义 —— 读共享 / 写独占的
// 判定一律由**真实的**进程内占用表给出（tool.NewFileClaimer），不拿自造的假表去测。

// claimProbeTimeout 是「第一路必须进入闸门 / 一路必须结束」这类断言的等待上限：足够宽
// 以免依赖机器速度，又足够短以免偶发死锁拖成整包超时。
const claimProbeTimeout = 2 * time.Second

// legResult 是一路调用的结果。
type legResult struct {
	val any
	err error
}

// recordingClaimer 包住真实的进程内占用表（tool.NewFileClaimer），把每一次登记尝试的目标
// 键记下来：测试因此能断言「哪些调用参与了互斥、此刻占用几个键」。
type recordingClaimer struct {
	inner tool.FileClaimer
	mu    sync.Mutex
	// live 镜像真实占用表的键集合：目标键 -> 路标识 -> 登记次数（同路重复登记按幂等算，
	// 与真实表一致；释放到最后一个占用者时删键）。
	live map[string]map[string]int
	// attempted 按发生顺序记录**每一次**登记尝试的目标键（含被判忙而失败的那些）：
	// 「尝试过」与「成功占到键」是两件事 —— 前者证明这一路确实拿着声明的目标来问过表。
	attempted []string
}

func newRecordingClaimer() *recordingClaimer {
	return &recordingClaimer{inner: tool.NewFileClaimer(), live: map[string]map[string]int{}}
}

// Claim 实现 tool.FileClaimer：先记账（判忙的尝试也要留痕），再让真实占用表判；只有
// 成功了才占键（失败即未登记，与真实表一致）。
func (r *recordingClaimer) Claim(ctx context.Context, legID, target string) error {
	r.mu.Lock()
	r.attempted = append(r.attempted, target)
	r.mu.Unlock()
	if err := r.inner.Claim(ctx, legID, target); err != nil {
		return err
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	legs := r.live[target]
	if legs == nil {
		legs = map[string]int{}
		r.live[target] = legs
	}
	legs[legID]++
	return nil
}

// Release 实现 tool.FileClaimer：无条件转发给真实占用表，并同步镜像（非占用者的释放是
// 空操作，与真实表一致）。
func (r *recordingClaimer) Release(ctx context.Context, legID, target string) {
	r.inner.Release(ctx, legID, target)
	r.mu.Lock()
	defer r.mu.Unlock()
	legs := r.live[target]
	if legs == nil {
		return
	}
	if n := legs[legID]; n > 1 {
		legs[legID] = n - 1
	} else {
		delete(legs, legID)
	}
	if len(legs) == 0 {
		delete(r.live, target)
	}
}

// liveCount 返回此刻占用的目标键数（= 真实占用表的键数）。
func (r *recordingClaimer) liveCount() int {
	r.mu.Lock()
	defer r.mu.Unlock()
	return len(r.live)
}

// claimAttempts 返回登记**尝试**的次数（未声明目标的调用必须为 0：它连表都不问）。
func (r *recordingClaimer) claimAttempts() int {
	r.mu.Lock()
	defer r.mu.Unlock()
	return len(r.attempted)
}

// targets 返回登记尝试过的目标键（副本）。
func (r *recordingClaimer) targets() []string {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := make([]string, len(r.attempted))
	copy(out, r.attempted)
	return out
}

// claimRig 是一套「真实插件 + 真实分发器」的互斥夹具：DataDir 即 file 工具唯一的 root，
// 占用表换成记录版。
type claimRig struct {
	d     *tool.Dispatcher
	rec   *recordingClaimer
	dir   string
	tools map[string]sdk.Tool
}

func newClaimRig(t *testing.T) *claimRig {
	t.Helper()
	dir := t.TempDir()
	rec := newRecordingClaimer()
	d, err := tool.New(registry.New(), tool.WithFileClaimer(rec))
	if err != nil {
		t.Fatalf("tool.New: %v", err)
	}
	set := buildTools(t, &config.Config{DataDir: dir})
	byName := make(map[string]sdk.Tool, len(set))
	for _, tk := range set {
		byName[tk.Name()] = tk
		if err := d.Register(tk); err != nil {
			t.Fatalf("Register %q: %v", tk.Name(), err)
		}
	}
	return &claimRig{d: d, rec: rec, dir: dir, tools: byName}
}

// seed 在沙箱 root 里放一个文件（测试夹具，不经过工具）。
func (r *claimRig) seed(t *testing.T, name, content string) string {
	t.Helper()
	p := filepath.Join(r.dir, name)
	if err := os.WriteFile(p, []byte(content), 0o644); err != nil {
		t.Fatalf("seed %q: %v", name, err)
	}
	return p
}

// gated 注册一个带闸门的**真实插件**实例（名字换成 name，声明面原样转发）。
func (r *claimRig) gated(t *testing.T, toolName, name string) *gatedTool {
	t.Helper()
	inner, ok := r.tools[toolName]
	if !ok {
		t.Fatalf("工具集里没有 %q", toolName)
	}
	g := newGatedTool(inner, name)
	if err := r.d.Register(g); err != nil {
		t.Fatalf("Register %q: %v", name, err)
	}
	return g
}

// gatedTool 包住一个真实工具，在 Execute 前后加两个可控闸门（进入信号 + 让行等待）：占用
// 登记在 Execute **之前**、释放在它**之后**，所以停在闸门里就等于「这一路持有声明出来的
// 占用」，撞车因此是确定性事件，而不是靠 sleep 抢时序。它还必须显式转发 FileTargets ——
// 内嵌接口只会提升接口里的方法，声明面不转发就变成了「未实现 FileTargeter 的工具」，
// 用例会退化成一个假绿。
type gatedTool struct {
	inner      sdk.Tool
	name       string
	entered    chan struct{}
	hold       chan struct{}
	enteredOne sync.Once
	holdOne    sync.Once
}

func newGatedTool(inner sdk.Tool, name string) *gatedTool {
	return &gatedTool{inner: inner, name: name, entered: make(chan struct{}), hold: make(chan struct{})}
}

func (g *gatedTool) Name() string { return g.name }

func (g *gatedTool) Description() string { return g.inner.Description() }

func (g *gatedTool) Parameters() (map[string]any, error) { return g.inner.Parameters() }

// FileTargets 原样转发内层工具的声明：测试验证的始终是**插件自己**的目标判定。
func (g *gatedTool) FileTargets(args any) []tool.FileTarget {
	targeter, ok := g.inner.(tool.FileTargeter)
	if !ok {
		return nil
	}
	return targeter.FileTargets(args)
}

func (g *gatedTool) Execute(ctx context.Context, args any) (any, error) {
	g.enteredOne.Do(func() { close(g.entered) })
	select {
	case <-g.hold:
	case <-ctx.Done():
		return nil, ctx.Err()
	}
	return g.inner.Execute(ctx, args)
}

// release 放行闸门（幂等：不会二次 close）。
func (g *gatedTool) release() { g.holdOne.Do(func() { close(g.hold) }) }

// invokeLeg 在指定并发路上同步调用工具。
func invokeLeg(d *tool.Dispatcher, leg, name string, args map[string]any) (any, error) {
	return d.Invoke(tool.WithLegID(context.Background(), leg), name, args)
}

// runLeg 在指定并发路上**异步**调用工具：它一进入闸门，持有该目标的占用就已经成立。
func runLeg(d *tool.Dispatcher, leg, name string, args map[string]any) <-chan legResult {
	out := make(chan legResult, 1)
	go func() {
		val, err := invokeLeg(d, leg, name, args)
		out <- legResult{val: val, err: err}
	}()
	return out
}

// awaitEnter 等一路真的进入 Execute（此刻它的占用已登记）。
func awaitEnter(t *testing.T, ch <-chan struct{}) {
	t.Helper()
	select {
	case <-ch:
	case <-time.After(claimProbeTimeout):
		t.Fatal("第一路没有进入执行")
	}
}

// awaitDone 等一路结束。
func awaitDone(t *testing.T, ch <-chan legResult) legResult {
	t.Helper()
	select {
	case res := <-ch:
		return res
	case <-time.After(claimProbeTimeout):
		t.Fatal("这一路没有结束")
		return legResult{}
	}
}

// assertBusy 断言一次调用**被判忙**：errors.Is 命中 tool.ErrFileBusy，且面向模型的文案里
// 带文件名与文件状态（说出是谁在动它、本次做不了什么）。
func assertBusy(t *testing.T, err error, wantSubs ...string) {
	t.Helper()
	if err == nil {
		t.Fatal("期望被判忙，实际调用成功了")
	}
	if !errors.Is(err, tool.ErrFileBusy) {
		t.Fatalf("错误不是 ErrFileBusy: %v", err)
	}
	msg := err.Error()
	for _, want := range append([]string{"请稍后再试"}, wantSubs...) {
		if !strings.Contains(msg, want) {
			t.Fatalf("回灌给模型的文案缺少 %q: %s", want, msg)
		}
	}
}

// assertNotBusy 断言一次调用没被判忙（其它失败允许：这里只关心互斥不该拦它）。
func assertNotBusy(t *testing.T, err error) {
	t.Helper()
	if errors.Is(err, tool.ErrFileBusy) {
		t.Fatalf("不该被判忙，得到 %v", err)
	}
}

// assertNoClaims 断言此刻占用表为空、且登记**尝试**次数恰为 wantAttempts：未声明目标的
// 调用必须一次都不问表，声明了但被判忙的那些则只留尝试、不占键。
func assertNoClaims(t *testing.T, rec *recordingClaimer, wantAttempts int) {
	t.Helper()
	if got := rec.liveCount(); got != 0 {
		t.Fatalf("占用表必须为空，残留 %d 个键", got)
	}
	if got := rec.claimAttempts(); got != wantAttempts {
		t.Fatalf("登记尝试次数 = %d，期望 %d", got, wantAttempts)
	}
}

// plainTool 是最小 sdk.Tool：**只**实现 sdk.Tool（不实现 FileTargeter），用来钉住「未声明
// 目标的工具完全不参与互斥」。
type plainTool struct{ name string }

func (p *plainTool) Name() string { return p.name }

func (p *plainTool) Description() string { return "test tool without file targets" }

func (p *plainTool) Parameters() (map[string]any, error) {
	return map[string]any{
		"type":       "object",
		"properties": map[string]any{"echo": map[string]any{"type": "string"}},
	}, nil
}

func (p *plainTool) Execute(_ context.Context, args any) (any, error) {
	m, _ := args.(map[string]any)
	return argString(m, "echo"), nil
}

// toolContent 取 file 工具读到的正文。
func toolContent(t *testing.T, val any) string {
	t.Helper()
	m, ok := val.(map[string]any)
	if !ok {
		t.Fatalf("读取结果不是对象: %#v", val)
	}
	content, _ := m["content"].(string)
	return content
}