Cyrene_v2/internal/runtime/loop/capacity_loop_test.go

package loop

import (
	"context"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/capacity"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/tool"
)

// capacityToolSet 为验证允许集过滤所用的四个代表性工具
// 构建一个 dispatcher：两个 quiet 工具、一个仅 awake 的，
// 以及一个在 quiet/awake 下永不被允许的。
func capacityToolSet(t *testing.T) *tool.Dispatcher {
	t.Helper()
	return dispatcherWith(t,
		&mockTool{name: "terminal"},
		&mockTool{name: "file"},
		&mockTool{name: "web-search"},
		&mockTool{name: "json"},
	)
}

// toolNames 返回 specs 中广告出的工具名集合。registry 的
// List 遍历 map，因此注册顺序不保留；做与顺序无关的断言时
// 应与集合（经 toolSet）比较，而不是与切片比较。
func toolNames(specs []model.ToolSpec) []string {
	out := make([]string, 0, len(specs))
	for _, s := range specs {
		out = append(out, s.Name)
	}
	return out
}

func toolSet(specs []model.ToolSpec) map[string]bool {
	set := make(map[string]bool, len(specs))
	for _, s := range specs {
		set[s.Name] = true
	}
	return set
}

func hasAll(set map[string]bool, names ...string) bool {
	for _, n := range names {
		if !set[n] {
			return false
		}
	}
	return true
}

// TestLoop_Capacity_QuietHidesHeavierTools 验证启用 capacity
// 判定器时，闲聊 turn 在流式路径上只向模型广告 quiet 允许集
// （terminal + file，不含 web-search/json）。
func TestLoop_Capacity_QuietHidesHeavierTools(t *testing.T) {
	t.Parallel()
	td := capacityToolSet(t)
	ad := &mockAdapter{stream: []mgResp{{resp: model.Response{Content: "陪你"}}}}
	a := newTestAgent(t, ad, td)
	dec := capacity.New(capacity.Config{Enabled: true, MinConfidence: 0.6})
	l := New(Options{
		Agent:    a,
		MaxRounds: 8,
		OnStream: func(model.StreamEvent) {},
		Capacity: dec,
	})
	if _, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "你好呀，抱抱你")}); err != nil {
		t.Fatalf("Run: %v", err)
	}
	if len(ad.streamReqs) == 0 {
		t.Fatal("no streaming request captured")
	}
	set := toolSet(ad.streamReqs[0].Tools)
	if !hasAll(set, "terminal", "file") {
		t.Fatalf("quiet tools = %v, want to include terminal+file", toolNames(ad.streamReqs[0].Tools))
	}
	if set["web-search"] || set["json"] {
		t.Fatalf("quiet tools = %v, want to exclude web-search/json", toolNames(ad.streamReqs[0].Tools))
	}
}

// TestLoop_Capacity_WakeExposesWebSearch 验证任务 turn 向模型
// 暴露 awake 允许集（增加 web-search）。
func TestLoop_Capacity_WakeExposesWebSearch(t *testing.T) {
	t.Parallel()
	td := capacityToolSet(t)
	ad := &mockAdapter{stream: []mgResp{{resp: model.Response{Content: "好的，这就查"}}}}
	a := newTestAgent(t, ad, td)
	dec := capacity.New(capacity.Config{Enabled: true, MinConfidence: 0.6})
	l := New(Options{
		Agent:    a,
		MaxRounds: 8,
		OnStream: func(model.StreamEvent) {},
		Capacity: dec,
	})
	if _, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "帮我查一下最近的天气")}); err != nil {
		t.Fatalf("Run: %v", err)
	}
	set := toolSet(ad.streamReqs[0].Tools)
	if !hasAll(set, "terminal", "file", "web-search") {
		t.Fatalf("awake tools = %v, want to include terminal+file+web-search", toolNames(ad.streamReqs[0].Tools))
	}
	if set["json"] {
		t.Fatalf("awake tools = %v, want to exclude json", toolNames(ad.streamReqs[0].Tools))
	}
}

// TestLoop_Capacity_NilKeepsAllTools 验证既有行为：没有
// capacity 判定器时，每个已注册工具都被广告（不过滤）。
func TestLoop_Capacity_NilKeepsAllTools(t *testing.T) {
	t.Parallel()
	td := capacityToolSet(t)
	ad := &mockAdapter{stream: []mgResp{{resp: model.Response{Content: "你好"}}}}
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8, OnStream: func(model.StreamEvent) {}})
	if _, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "你好")}); err != nil {
		t.Fatalf("Run: %v", err)
	}
	set := toolSet(ad.streamReqs[0].Tools)
	if !hasAll(set, "terminal", "file", "web-search", "json") {
		t.Fatalf("nil capacity tools = %v, want all four", toolNames(ad.streamReqs[0].Tools))
	}
	if len(set) != 4 {
		t.Fatalf("nil capacity tools = %v, want exactly four", toolNames(ad.streamReqs[0].Tools))
	}
}

// TestLoop_Capacity_EmitsAwakenEvent 验证状态变化在流式路径上
// 经 OnCapacity 钩子浮出一个 capacity_awaken 事件。
func TestLoop_Capacity_EmitsAwakenEvent(t *testing.T) {
	t.Parallel()
	td := capacityToolSet(t)
	ad := &mockAdapter{stream: []mgResp{{resp: model.Response{Content: "查到了"}}}}
	a := newTestAgent(t, ad, td)
	dec := capacity.New(capacity.Config{Enabled: true, MinConfidence: 0.6})
	var got []capacity.Event
	l := New(Options{
		Agent:      a,
		MaxRounds:  8,
		OnStream:   func(model.StreamEvent) {},
		Capacity:   dec,
		OnCapacity: func(ev capacity.Event) { got = append(got, ev) },
	})
	if _, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "帮我查一下最近的天气")}); err != nil {
		t.Fatalf("Run: %v", err)
	}
	if len(got) != 1 {
		t.Fatalf("capacity events = %d, want 1 (got %+v)", len(got), got)
	}
	ev := got[0]
	if ev.Type != capacity.EventAwaken || ev.FromState != capacity.StateQuiet || ev.ToState != capacity.StateAwake {
		t.Fatalf("event = %+v, want quiet→awake awaken", ev)
	}
	if ev.Tool != "web-search" {
		t.Fatalf("lit tool = %q, want web-search", ev.Tool)
	}
}