Cyrene_v2/internal/runtime/capacity/capacity_test.go

package capacity

import (
	"context"
	"reflect"
	"testing"
)

// allTools is a representative registration-order tool list for tests; it covers
// the basic/io/terminal plugins actually shipped by the repo.
func allTools() []string {
	return []string{
		"terminal", "file", "knowledge", "text",
		"web-search", "web-fetch", "http",
		"calculator", "datetime", "random", "json", "markdown",
	}
}

func activeConfig() Config {
	return Config{Enabled: true, MinConfidence: 0.6}
}

func in(messages ...string) Input {
	return Input{Messages: messages, CurrentState: StateQuiet, AllToolNames: allTools()}
}

func contains(s []string, want string) bool {
	for _, v := range s {
		if v == want {
			return true
		}
	}
	return false
}

func TestDecide_Disabled_ReturnsAllToolsNoEvents(t *testing.T) {
	d := New(Config{}) // Enabled zero → disabled
	dec := d.Decide(context.Background(), in("帮我查一下资料"))
	if !reflect.DeepEqual(dec.AllowedTools, allTools()) {
		t.Fatalf("disabled allowed = %v, want all %v", dec.AllowedTools, allTools())
	}
	if len(dec.Events) != 0 {
		t.Fatalf("disabled should emit no events, got %+v", dec.Events)
	}
}

func TestDecide_Chitchat_StaysQuiet(t *testing.T) {
	d := New(activeConfig())
	for _, msg := range []string{"你好", "抱抱", "嗯嗯", "晚安", "今天好累呀"} {
		if got := d.Decide(context.Background(), in(msg)); got.State != StateQuiet {
			t.Errorf("chitchat %q woke to %s, want quiet", msg, got.State)
		}
	}
}

func TestDecide_Chitchat_NoAwakenEvent(t *testing.T) {
	d := New(activeConfig())
	dec := d.Decide(context.Background(), in("你好"))
	if len(dec.Events) != 0 {
		t.Fatalf("quiet chat should not emit capacity_awaken, got %+v", dec.Events)
	}
}

func TestDecide_Query_Awakens(t *testing.T) {
	d := New(activeConfig())
	dec := d.Decide(context.Background(), in("帮我查一下今天的天气"))
	if dec.State != StateAwake {
		t.Fatalf("query state = %s, want awake", dec.State)
	}
	if dec.Confidence < 0.6 {
		t.Fatalf("query confidence = %v, want >= 0.6", dec.Confidence)
	}
	if !contains(dec.AllowedTools, "web-search") {
		t.Fatalf("awake allowed should include web-search, got %v", dec.AllowedTools)
	}
	if len(dec.Events) != 1 || dec.Events[0].FromState != StateQuiet || dec.Events[0].ToState != StateAwake {
		t.Fatalf("expected a quiet→awake event, got %+v", dec.Events)
	}
}

func TestDecide_EditFile_Awakens(t *testing.T) {
	d := New(activeConfig())
	dec := d.Decide(context.Background(), in("改一下 cx.go 这个函数"))
	if dec.State != StateAwake {
		t.Fatalf("edit state = %s, want awake", dec.State)
	}
	if !contains(dec.AllowedTools, "file") {
		t.Fatalf("awake allowed should include file, got %v", dec.AllowedTools)
	}
	if len(dec.Events) != 1 || dec.Events[0].Tool != "file" {
		t.Fatalf("edit event should light up file, got %+v", dec.Events)
	}
}

func TestDecide_ComplexTask_Full(t *testing.T) {
	d := New(activeConfig())
	dec := d.Decide(context.Background(), in("帮我重构整个项目，迁移到新架构"))
	if dec.State != StateFull {
		t.Fatalf("complex state = %s, want full", dec.State)
	}
	if !reflect.DeepEqual(dec.AllowedTools, allTools()) {
		t.Fatalf("full allowed = %v, want all %v", dec.AllowedTools, allTools())
	}
}

func TestDecide_L1_UndecidedDeferstoL2_WakesOnTaskTrend(t *testing.T) {
	d := New(activeConfig())
	// Last turn (这个) is undecided by L1; the recent history is task-leaning.
	dec := d.Decide(context.Background(), in("帮我查一下资料", "还有吗", "这个"))
	if dec.State != StateAwake {
		t.Fatalf("task-trend state = %s, want awake", dec.State)
	}
}

func TestDecide_L2_CompanionTrend_Quiet(t *testing.T) {
	d := New(activeConfig())
	dec := d.Decide(context.Background(), in("你好", "抱抱", "嗯嗯", "这个"))
	if dec.State != StateQuiet {
		t.Fatalf("companion-trend state = %s, want quiet", dec.State)
	}
}

func TestDecide_L2_Followup_Wakes(t *testing.T) {
	d := New(activeConfig())
	// A clear previous task plus a bare continuation cue (继续) stays awake —
	// L2 picks up the follow-up on the previous round's task (design §2.2).
	dec := d.Decide(context.Background(), in("帮我查一下资料", "继续"))
	if dec.State != StateAwake {
		t.Fatalf("followup state = %s, want awake", dec.State)
	}
}

func TestDecide_Undecided_ConservativeQuiet(t *testing.T) {
	d := New(activeConfig())
	// Entirely neutral turns with no signal at all.
	dec := d.Decide(context.Background(), in("嗯", "这个", "然后"))
	if dec.State != StateQuiet {
		t.Fatalf("undecided state = %s, want quiet (conservative)", dec.State)
	}
}

func TestDecide_PersonaForceState(t *testing.T) {
	d := New(activeConfig())
	inp := in("你好")
	inp.Awakening = Awakening{ForceState: StateAwake}
	dec := d.Decide(context.Background(), inp)
	if dec.State != StateAwake {
		t.Fatalf("force_state state = %s, want awake", dec.State)
	}
}

func TestDecide_PersonaDefaultState(t *testing.T) {
	// An assistant-type persona rests at awake; an undecided round lands there.
	d := New(Config{Enabled: true, MinConfidence: 0.6, Awakening: Awakening{DefaultState: StateAwake}})
	dec := d.Decide(context.Background(), in("这个"))
	if dec.State != StateAwake {
		t.Fatalf("assistant default_state = %s, want awake", dec.State)
	}
}

func TestDecide_PersonaMinConfidence_ForcesQuiet(t *testing.T) {
	// A high threshold forces even a strong task signal to stay quiet. Low conf.
	d := New(Config{Enabled: true, MinConfidence: 0.6})
	inp := in("帮我查一下资料")
	inp.Awakening = Awakening{MinConfidence: 0.99}
	if dec := d.Decide(context.Background(), inp); dec.State != StateQuiet {
		t.Fatalf("high min_confidence should force quiet, got %s", dec.State)
	}
	// The default threshold lets the same strong signal wake.
	if dec := d.Decide(context.Background(), in("帮我查一下资料")); dec.State != StateAwake {
		t.Fatalf("default min_confidence should wake, got %s", dec.State)
	}
}

func TestDecide_AllowedSets(t *testing.T) {
	d := New(activeConfig())

	quiet := d.Decide(context.Background(), in("你好")).AllowedTools
	if !reflect.DeepEqual(quiet, []string{"terminal", "file", "knowledge", "text"}) {
		t.Fatalf("quiet allowed = %v", quiet)
	}

	awake := d.Decide(context.Background(), in("帮我查一下资料")).AllowedTools
	wantAwake := []string{"terminal", "file", "knowledge", "text", "web-search", "web-fetch", "http"}
	if !reflect.DeepEqual(awake, wantAwake) {
		t.Fatalf("awake allowed = %v, want %v", awake, wantAwake)
	}
	if contains(awake, "json") || contains(awake, "markdown") || contains(awake, "calculator") {
		t.Fatalf("awake must not expose heavy utility tools, got %v", awake)
	}

	full := d.Decide(context.Background(), in("帮我重构整个项目")).AllowedTools
	if !reflect.DeepEqual(full, allTools()) {
		t.Fatalf("full allowed = %v, want all %v", full, allTools())
	}
}

func TestDecide_Event_LitToolIsAllowed(t *testing.T) {
	d := New(activeConfig())
	dec := d.Decide(context.Background(), in("帮我查一下资料"))
	if len(dec.Events) != 1 {
		t.Fatalf("expected one event, got %+v", dec.Events)
	}
	ev := dec.Events[0]
	if !ev.IsAwaken() || ev.Tool == "" {
		t.Fatalf("event should be an awaken with a lit tool, got %+v", ev)
	}
	if !contains(dec.AllowedTools, ev.Tool) {
		t.Fatalf("lit tool %q not in allowed set %v", ev.Tool, dec.AllowedTools)
	}
}

func TestDecide_CustomToolSets(t *testing.T) {
	d := New(Config{
		Enabled:       true,
		MinConfidence: 0.6,
		ToolSets:      map[State][]string{StateQuiet: {"terminal"}},
	})
	dec := d.Decide(context.Background(), in("你好"))
	if !reflect.DeepEqual(dec.AllowedTools, []string{"terminal"}) {
		t.Fatalf("custom quiet allowed = %v, want [terminal]", dec.AllowedTools)
	}
}

func TestDecide_NilToolNames_EmptyAllowed(t *testing.T) {
	d := New(activeConfig())
	dec := d.Decide(context.Background(), Input{Messages: []string{"你好"}, CurrentState: StateQuiet})
	if dec.State != StateQuiet {
		t.Fatalf("state = %s, want quiet", dec.State)
	}
	if len(dec.AllowedTools) != 0 {
		t.Fatalf("allowed with no known tools = %v, want empty", dec.AllowedTools)
	}
}