Cyrene_v2/internal/runtime/orchestration/spawnwhitelist_test.go

package orchestration

import (
	"context"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/tool"
)

// twoToolAdapter 先在一轮里同时请求两个工具调用，见到任意工具结果后给最终
// 答案——使一次子代理运行同时覆盖「白名单内被放行」与「白名单外被否决」
// 两条路径。它同时捕获每次模型请求，供断言工具广告面与回灌文本。
type twoToolAdapter struct {
	first, second string
	final         string

	mu       sync.Mutex
	requests []model.Request
}

func (a *twoToolAdapter) Name() string { return "two" }

func (a *twoToolAdapter) Chat(_ context.Context, req model.Request, _ model.Options) (model.Response, error) {
	a.mu.Lock()
	a.requests = append(a.requests, req)
	a.mu.Unlock()
	for _, m := range req.Messages {
		if m.Role == model.RoleTool {
			return model.Response{Model: "mock-1", Content: a.final, StopReason: "end_turn"}, nil
		}
	}
	return model.Response{Model: "mock-1", ToolCalls: []model.ToolCall{
		{ID: "call-1", Name: a.first, Args: map[string]any{"n": 1}},
		{ID: "call-2", Name: a.second, Args: map[string]any{"n": 2}},
	}}, nil
}

func (a *twoToolAdapter) ChatStream(ctx context.Context, req model.Request, on func(model.StreamEvent)) (model.Response, error) {
	return a.Chat(ctx, req, model.Options{})
}

// request 返回第 i 次捕获的模型请求。
func (a *twoToolAdapter) request(i int) model.Request {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.requests[i]
}

func (a *twoToolAdapter) count() int {
	a.mu.Lock()
	defer a.mu.Unlock()
	return len(a.requests)
}

// hitTool 是一个只计数的工具（原子递增，容忍并行批次），使测试能断言
// 「真的执行了」而不仅是被广告。
type hitTool struct {
	name  string
	hits  *int32
	delay time.Duration
}

func (h *hitTool) Name() string                        { return h.name }
func (h *hitTool) Description() string                 { return "count its invocations" }
func (h *hitTool) Parameters() (map[string]any, error) { return nil, nil }

func (h *hitTool) Execute(ctx context.Context, _ any) (any, error) {
	if h.delay > 0 {
		select {
		case <-time.After(h.delay):
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
	if h.hits != nil {
		atomic.AddInt32(h.hits, 1)
	}
	return h.name + ":ok", nil
}

// twoToolRig 组装「两个可数工具 + 请求两工具并捕获请求的适配器 + 派发工具」
// 的测试台。
func twoToolRig(t *testing.T) (*SpawnTool, *twoToolAdapter, *int32, *int32) {
	t.Helper()
	var alphaHits, betaHits int32
	d, err := tool.New(registry.New())
	if err != nil {
		t.Fatalf("tool.New: %v", err)
	}
	for _, h := range []*hitTool{
		{name: "alpha", hits: &alphaHits},
		{name: "beta", hits: &betaHits},
	} {
		if rerr := d.Register(h); rerr != nil {
			t.Fatalf("Register: %v", rerr)
		}
	}
	two := &twoToolAdapter{first: "alpha", second: "beta", final: "完成"}
	s := captureLoop(t, two, d)
	sp := NewSpawnTool(s, WithSpawnWait(3*time.Second))
	return sp, two, &alphaHits, &betaHits
}

// TestToolWhitelistDefaultAllTools 断言缺省全开：不列出 tool_whitelist 的
// 派发，子代理照常执行主 agent 全量工具集里的每个工具（02.4.1）。
func TestToolWhitelistDefaultAllTools(t *testing.T) {
	t.Parallel()
	sp, ad, alphaHits, betaHits := twoToolRig(t)

	raw, err := sp.Execute(context.Background(), map[string]any{"goal": "both"})
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	if got := decodeSpawnResult(t, raw).Status; got != "done" {
		t.Fatalf("Status = %q, want done", got)
	}
	if *alphaHits == 0 || *betaHits == 0 {
		t.Fatalf("default dispatch restricted tools: alpha=%d beta=%d, want both executed", *alphaHits, *betaHits)
	}
	// 两轮请求都广告了全部工具（能力面未被收缩）。
	for i := 0; i < ad.count(); i++ {
		names := toolNamesOf(ad.request(i))
		if !names["alpha"] || !names["beta"] {
			t.Fatalf("request %d advertised tools %v, want the full set (default = inherit everything)", i, names)
		}
	}
}

// TestToolWhitelistExplicitRestricts 断言显式列出 = 只放行所列（02.4.1）：
// 白名单外的调用被否决为软结果回灌（模型可见 not_allowed 提示），白名单内
// 的工具照常执行。
func TestToolWhitelistExplicitRestricts(t *testing.T) {
	t.Parallel()
	sp, ad, alphaHits, betaHits := twoToolRig(t)
	raw, err := sp.Execute(context.Background(), map[string]any{
		"goal":           "only beta",
		"tool_whitelist": []any{"beta"},
	})
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	if got := decodeSpawnResult(t, raw).Status; got != "done" {
		t.Fatalf("Status = %q, want done (a veto is a soft result, not a failure)", got)
	}
	if *alphaHits != 0 {
		t.Fatalf("alpha executed %d times despite not being whitelisted", *alphaHits)
	}
	if *betaHits == 0 {
		t.Fatal("beta never executed although it is whitelisted")
	}
	// 被否决的调用以软结果回灌，模型能看到可读原因。
	vetoed := false
	for i := 0; i < ad.count(); i++ {
		for _, m := range ad.request(i).Messages {
			if m.Role == model.RoleTool && strings.Contains(messageTextOf(m), "not in this sub-agent's tool whitelist") {
				vetoed = true
			}
		}
	}
	if !vetoed {
		t.Fatal("the out-of-whitelist call never surfaced a readable veto to the model")
	}
}

// toolNamesOf 提取一次请求里广告的工具名集合。
func toolNamesOf(req model.Request) map[string]bool {
	out := make(map[string]bool, len(req.Tools))
	for _, ts := range req.Tools {
		out[ts.Name] = true
	}
	return out
}

// ctxCaptureRunner 记录运行 context，供断言白名单「随运行上下文传递」。
type ctxCaptureRunner struct {
	mu  sync.Mutex
	ctx []context.Context
}

func (r *ctxCaptureRunner) Run(ctx context.Context, _ string, _ []model.Message) (*loop.Result, error) {
	r.mu.Lock()
	r.ctx = append(r.ctx, ctx)
	r.mu.Unlock()
	return &loop.Result{Text: "ok", Steps: 1, StopReason: "stop"}, nil
}

func (r *ctxCaptureRunner) last() context.Context {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.ctx[len(r.ctx)-1]
}

// TestToolWhitelistRidesRunContext 断言名单通过 loop.WithToolWhitelist
// 挂到运行 context 上（空名单 = 未挂载 = 全开）。
func TestToolWhitelistRidesRunContext(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 1})
	m := &ctxCaptureRunner{}
	s.runner = m

	sub, err := s.Spawn(context.Background(), AgentSpec{Goal: "w", ToolWhitelist: []string{"a", "b"}})
	if err != nil {
		t.Fatalf("Spawn: %v", err)
	}
	if _, werr := s.Wait(context.Background(), []string{sub.ID}); werr != nil {
		t.Fatalf("Wait: %v", werr)
	}
	if got := loop.ToolWhitelistFrom(m.last()); len(got) != 2 || got[0] != "a" || got[1] != "b" {
		t.Fatalf("run context whitelist = %v, want [a b]", got)
	}

	bare, err := s.Spawn(context.Background(), AgentSpec{Goal: "bare"})
	if err != nil {
		t.Fatalf("Spawn bare: %v", err)
	}
	if _, werr := s.Wait(context.Background(), []string{bare.ID}); werr != nil {
		t.Fatalf("Wait bare: %v", werr)
	}
	if got := loop.ToolWhitelistFrom(m.last()); got != nil {
		t.Fatalf("whitelist = %v on a whitelist-less dispatch, want nil (全开)", got)
	}
}