sundynix-agentix/sundynix-dispatcher/internal/llm/failover_test.go

package llm

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/cloudwego/eino/components/model"
	"github.com/cloudwego/eino/schema"

	"github.com/sundynix/sundynix-dispatcher/internal/harness"
)

// fakeModel 是 model.ToolCallingChatModel 测试替身：fail=true 则 Generate/Stream 报错。
type fakeModel struct {
	name  string
	reply string
	fail  bool
	calls *int // 共享计数器，记录被调次数（验证 failover 没多调/早停）
}

func (f *fakeModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
	if f.calls != nil {
		*f.calls++
	}
	if f.fail {
		return nil, fmt.Errorf("%s down", f.name)
	}
	return schema.AssistantMessage(f.reply, nil), nil
}

func (f *fakeModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
	if f.calls != nil {
		*f.calls++
	}
	if f.fail {
		return nil, fmt.Errorf("%s down", f.name)
	}
	sr, sw := schema.Pipe[*schema.Message](1)
	go func() { sw.Send(schema.AssistantMessage(f.reply, nil), nil); sw.Close() }()
	return sr, nil
}

func (f *fakeModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { return f, nil }

func genText(t *testing.T, ctx context.Context, m model.ToolCallingChatModel) (string, error) {
	t.Helper()
	out, err := m.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
	if err != nil {
		return "", err
	}
	return out.Content, nil
}

// TestFailover_PrimaryWorks 主可用 → 直接用主，备用不被调用。
func TestFailover_PrimaryWorks(t *testing.T) {
	var pc, fc int
	m := newFailoverModel([]model.ToolCallingChatModel{
		&fakeModel{name: "primary", reply: "主回答", calls: &pc},
		&fakeModel{name: "fb", reply: "备回答", calls: &fc},
	}, nil)
	ans, err := genText(t, context.Background(), m)
	if err != nil || ans != "主回答" {
		t.Fatalf("应返回主回答, got %q err=%v", ans, err)
	}
	if pc != 1 || fc != 0 {
		t.Fatalf("主可用时不应调备用: primary=%d fb=%d", pc, fc)
	}
}

// TestFailover_PrimaryDownUseFallback 主挂 → 切备用并成功。
func TestFailover_PrimaryDownUseFallback(t *testing.T) {
	var pc, fc int
	switched := false
	m := newFailoverModel([]model.ToolCallingChatModel{
		&fakeModel{name: "primary", fail: true, calls: &pc},
		&fakeModel{name: "fb", reply: "备回答", calls: &fc},
	}, func(int, error) { switched = true })
	ans, err := genText(t, context.Background(), m)
	if err != nil || ans != "备回答" {
		t.Fatalf("主挂应切备用, got %q err=%v", ans, err)
	}
	if pc != 1 || fc != 1 || !switched {
		t.Fatalf("主备各调一次且触发回调: primary=%d fb=%d switched=%v", pc, fc, switched)
	}
}

// TestFailover_AllDown 全挂 → 返回最后一个错误。
func TestFailover_AllDown(t *testing.T) {
	m := newFailoverModel([]model.ToolCallingChatModel{
		&fakeModel{name: "p", fail: true},
		&fakeModel{name: "fb", fail: true},
	}, nil)
	if _, err := genText(t, context.Background(), m); err == nil {
		t.Fatal("全挂应返回错误")
	}
}

// TestFailover_CtxCanceledNoSwitch 调用方已取消 → 不切备用（切了也没用，且违背意图）。
func TestFailover_CtxCanceledNoSwitch(t *testing.T) {
	var pc, fc int
	ctx, cancel := context.WithCancel(context.Background())
	cancel() // 立即取消
	m := newFailoverModel([]model.ToolCallingChatModel{
		&fakeModel{name: "p", fail: true, calls: &pc},
		&fakeModel{name: "fb", reply: "备", calls: &fc},
	}, nil)
	if _, err := genText(t, ctx, m); err == nil {
		t.Fatal("已取消仍应返回主错误")
	}
	if fc != 0 {
		t.Fatalf("取消时不应切备用: fb 被调 %d 次", fc)
	}
}

// TestFailover_StreamFailover Stream 建流失败也切备用。
func TestFailover_StreamFailover(t *testing.T) {
	m := newFailoverModel([]model.ToolCallingChatModel{
		&fakeModel{name: "p", fail: true},
		&fakeModel{name: "fb", reply: "流备"},
	}, nil)
	sr, err := m.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")})
	if err != nil {
		t.Fatalf("Stream 应切到备用: %v", err)
	}
	defer sr.Close()
	chunk, _ := sr.Recv()
	if chunk == nil || chunk.Content != "流备" {
		t.Fatalf("应收到备用的流, got %+v", chunk)
	}
}

// 主模型持续失败达阈值 → 熔断 → 后续请求跳过主、直连备用（省掉每次白试主的失败往返）。
func TestFailover_TrippedPrimarySkipped(t *testing.T) {
	pc, fc := 0, 0
	p := &fakeModel{name: "p", fail: true, calls: &pc}
	fb := &fakeModel{name: "fb", reply: "备", calls: &fc}
	m := newFailoverModel([]model.ToolCallingChatModel{p, fb}, nil)
	// 连打 threshold+2 次：前 threshold 次主被调并失败→切备；熔断后主被跳过。
	for i := 0; i < fbBreakerThreshold+2; i++ {
		if ans, err := genText(t, context.Background(), m); err != nil || ans != "备" {
			t.Fatalf("每次都应最终拿到备用回答, got %q err=%v", ans, err)
		}
	}
	if pc != fbBreakerThreshold {
		t.Fatalf("主应在失败 %d 次后被熔断跳过，实际被调 %d 次", fbBreakerThreshold, pc)
	}
	if fc != fbBreakerThreshold+2 {
		t.Fatalf("备用应每次都被调: got %d", fc)
	}
}

// 关键：WithTools 重包必须共享熔断状态。主已熔断 → 新链仍跳过主（否则 rewrap 清零 = 熔断失效）。
func TestFailover_WithToolsSharesBreakerState(t *testing.T) {
	pc, fc := 0, 0
	p := &fakeModel{name: "p", fail: true, calls: &pc}
	fb := &fakeModel{name: "fb", reply: "备", calls: &fc}
	m := newFailoverModel([]model.ToolCallingChatModel{p, fb}, nil)
	for i := 0; i < fbBreakerThreshold; i++ { // 打到主熔断
		_, _ = genText(t, context.Background(), m)
	}
	if pc != fbBreakerThreshold {
		t.Fatalf("主应被调 %d 次, got %d", fbBreakerThreshold, pc)
	}
	bound, err := m.WithTools(nil)
	if err != nil {
		t.Fatal(err)
	}
	before := pc
	if _, err := genText(t, context.Background(), bound); err != nil {
		t.Fatal(err)
	}
	if pc != before {
		t.Fatalf("重包后主不应再被调（熔断状态须共享），却多调了 %d 次", pc-before)
	}
}

// 熔断后冷却到点 → 半开探测打回主 → 主已恢复则成功切回主。
func TestFailover_RecoversAfterCooldown(t *testing.T) {
	pc, fc := 0, 0
	p := &fakeModel{name: "p", reply: "主回答", fail: true, calls: &pc}
	fb := &fakeModel{name: "fb", reply: "备", calls: &fc}
	m := &failoverModel{
		models: []model.ToolCallingChatModel{p, fb},
		breakers: []*harness.CircuitBreaker{
			harness.NewCircuitBreakerWith(fbBreakerThreshold, 30*time.Millisecond, 1),
			harness.NewCircuitBreakerWith(fbBreakerThreshold, 30*time.Millisecond, 1),
		},
	}
	for i := 0; i < fbBreakerThreshold; i++ { // 打到主熔断
		_, _ = genText(t, context.Background(), m)
	}
	tripped := pc
	p.fail = false // 主恢复健康
	// 冷却前：主仍被跳过。
	if _, _ = genText(t, context.Background(), m); pc != tripped {
		t.Fatalf("冷却前主应仍被跳过，却被调用了")
	}
	// 冷却到点：半开探测打回主 → 成功恢复。
	time.Sleep(40 * time.Millisecond)
	ans, err := genText(t, context.Background(), m)
	if err != nil || ans != "主回答" {
		t.Fatalf("冷却后应探测回主并成功, got %q err=%v", ans, err)
	}
	if pc <= tripped {
		t.Fatalf("冷却后主应被再次探测调用")
	}
}

// TestFailover_WithToolsBindsAll WithTools 给链上每个模型绑定，返回仍是 failover 链。
func TestFailover_WithToolsBindsAll(t *testing.T) {
	m := newFailoverModel([]model.ToolCallingChatModel{
		&fakeModel{name: "p", fail: true},
		&fakeModel{name: "fb", reply: "带工具的备"},
	}, nil)
	bound, err := m.WithTools(nil)
	if err != nil {
		t.Fatal(err)
	}
	ans, err := genText(t, context.Background(), bound)
	if err != nil || ans != "带工具的备" {
		t.Fatalf("绑定工具后 failover 仍生效, got %q err=%v", ans, err)
	}
}