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 builds a dispatcher over the four representative tools used to
// verify the allowed-set filter: two quiet tools, one awake-only, and one that is
// never allowed in 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 returns the set of tool names advertised in specs. The registry
// List iterates a map, so registration order is not preserved; compare with a
// set (via toolSet) rather than a slice for order-independent assertions.
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 verifies that with an enabled capacity
// decider a chitchat turn advertises only the quiet allowed set to the model
// (terminal + file, not web-search/json) on the streaming path.
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 verifies a task turn exposes the awake
// allowed set (adds web-search) to the model.
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 verifies the pre-existing behavior: without
// a capacity decider every registered tool is advertised (no filtering).
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 verifies a state change surfaces a
// capacity_awaken event through the OnCapacity hook on the streaming path.
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)
	}
}