Cyrene_v2/internal/runtime/orchestration/scheduler_test.go

package orchestration

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
	"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/orchestration/comm"
)

const waitTimeout = 3 * time.Second

// mockRunner 是供调度器测试使用的可配置 runner。它统计并发调用数（供并行上限测试）、
// 捕获收到的消息（供上下文隔离测试），并可在 gate channel 上阻塞，直到它被关闭或 ctx
// 结束（供中止/超时测试把一个子代理维持在「运行中」）。
type mockRunner struct {
	mu        sync.Mutex
	gate      chan struct{}
	active    int
	maxActive int
	total     int
	calls     [][]model.Message
	result    *loop.Result
	err       error
}

func (m *mockRunner) Run(ctx context.Context, _ string, msgs []model.Message) (*loop.Result, error) {
	m.mu.Lock()
	m.active++
	if m.active > m.maxActive {
		m.maxActive = m.active
	}
	m.total++
	m.calls = append(m.calls, cloneMessages(msgs))
	m.mu.Unlock()

	var gerr error
	if m.gate != nil {
		select {
		case <-m.gate:
		case <-ctx.Done():
			gerr = ctx.Err()
		}
	}

	m.mu.Lock()
	m.active--
	m.mu.Unlock()

	if gerr != nil {
		return nil, gerr
	}
	if m.err != nil {
		return nil, m.err
	}
	if m.result != nil {
		return m.result, nil
	}
	return &loop.Result{Text: "ok", Steps: 1, StopReason: "stop"}, nil
}

func (m *mockRunner) activeCount() int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.active
}

func (m *mockRunner) maxActiveCount() int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.maxActive
}

func (m *mockRunner) captured(i int) []model.Message {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.calls[i]
}

// mockAdapter 是一个最小 model.Adapter，使真实 loop.Loop 无需供应商即可在测试中驱动。
// Chat 返回一个脚本化的最终响应。
type mockAdapter struct {
	name string
	resp model.Response
	err  error
}

func (m *mockAdapter) Name() string {
	if m.name == "" {
		return "mock"
	}
	return m.name
}

func (m *mockAdapter) Chat(context.Context, model.Request, model.Options) (model.Response, error) {
	return m.resp, m.err
}

func (m *mockAdapter) ChatStream(context.Context, model.Request, func(model.StreamEvent)) (model.Response, error) {
	return m.resp, m.err
}

// newScheduler 在一个全新的进程内总线上构建 Scheduler，并接上默认 comm 探针。总线在
// 测试清理时关闭。调用方在调用之后再设置 sched.runner，或传 Options.Loop 驱动真实循环。
func newScheduler(t *testing.T, opts Options) *Scheduler {
	t.Helper()
	b := bus.New()
	t.Cleanup(func() { _ = b.Close() })
	opts.Bus = b
	return New(opts)
}

func waitFor(t *testing.T, cond func() bool, d time.Duration) {
	t.Helper()
	deadline := time.Now().Add(d)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatalf("condition not met within %v", d)
}

func countStatus(subs []*SubAgent, st Status) int {
	n := 0
	for _, s := range subs {
		if s.Status() == st {
			n++
		}
	}
	return n
}

// collectEventNames 从 ch 读取事件，直到出现终结态事件或期限到期，按顺序返回事件名。
func collectEventNames(t *testing.T, ch <-chan comm.Event, d time.Duration) []string {
	t.Helper()
	var names []string
	deadline := time.After(d)
	for {
		select {
		case e := <-ch:
			names = append(names, e.Name)
			if e.Name == EventDone || e.Name == EventFailed || e.Name == EventAborted {
				return names
			}
		case <-deadline:
			return names
		}
	}
}

func TestSpawnStatusFlow(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 2})
	m := &mockRunner{gate: make(chan struct{})}
	s.runner = m

	sub, err := s.Spawn(context.Background(), AgentSpec{Goal: "do X"})
	if err != nil {
		t.Fatalf("Spawn: %v", err)
	}
	if sub.ID == "" {
		t.Fatal("Spawn returned an empty id")
	}
	if sub.Goal != "do X" {
		t.Fatalf("Goal = %q, want do X", sub.Goal)
	}
	if sub.WorkingMemory == nil {
		t.Fatal("sub-agent has no isolated working memory")
	}

	// pending -> running -> done。
	waitFor(t, func() bool { return sub.Status() == StatusRunning }, waitTimeout)
	close(m.gate)
	waitFor(t, func() bool { return sub.Status() == StatusDone }, waitTimeout)

	if !sub.Done() {
		t.Fatal("sub-agent not terminal after completion")
	}
	if sub.Summary() == "" {
		t.Fatal("summary was not written on success")
	}
	if sub.Result() == nil {
		t.Fatal("result was not stored on success")
	}
	if sub.Err() != nil {
		t.Fatalf("Err = %v, want nil", sub.Err())
	}
	if sub.FinishedAt().IsZero() {
		t.Fatal("FinishedAt is not set on a terminal sub-agent")
	}
}

func TestSpawn_MissingGoal(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{})
	if _, err := s.Spawn(context.Background(), AgentSpec{}); !errors.Is(err, ErrMissingGoal) {
		t.Fatalf("Spawn err = %v, want ErrMissingGoal", err)
	}
}

func TestParallelLimit(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 2})
	gate := make(chan struct{})
	m := &mockRunner{gate: gate}
	s.runner = m

	ctx := context.Background()
	var subs []*SubAgent
	for i := 0; i < 5; i++ {
		sub, err := s.Spawn(ctx, AgentSpec{Goal: fmt.Sprintf("task %d", i)})
		if err != nil {
			t.Fatalf("Spawn %d: %v", i, err)
		}
		subs = append(subs, sub)
	}

	// 等到两个并行槽位都被占用，再确认并发运行数从未超过 MaxParallel。这里等的是
	// runner 内部的活跃计数，而不是「状态已翻成 running」：状态是在进入 runner 之前
	// 置位的，两者之间有一个调度窗口，用墙钟轮询去撞它必然偶发（仅看状态时读到过
	// maxActive = 0/1）。gate 未开，因此 active == 2 是一个稳定状态。
	waitFor(t, func() bool { return m.activeCount() >= 2 }, waitTimeout)
	if got := m.maxActiveCount(); got != 2 {
		t.Fatalf("max concurrent runs = %d, want 2 while gated", got)
	}
	if running := countStatus(subs, StatusRunning); running != 2 {
		t.Fatalf("running = %d, want 2", running)
	}
	if pending := countStatus(subs, StatusPending); pending != 3 {
		t.Fatalf("pending = %d, want 3 (queued behind the parallel limit)", pending)
	}

	close(gate)
	waitFor(t, func() bool { return countStatus(subs, StatusDone) == 5 }, waitTimeout)

	res, err := s.Join(ctx)
	if err != nil {
		t.Fatalf("Join: %v", err)
	}
	if res.Total != 5 || res.Succeeded != 5 {
		t.Fatalf("merged = %d total / %d succeeded, want 5/5", res.Total, res.Succeeded)
	}
	// 峰值在整轮排空中都不得越界：这条覆盖放行后逐个进出槽位的阶段，
	// 因此它也校验每个运行体结束时确实归还了槽位。
	if got := m.maxActiveCount(); got != 2 {
		t.Fatalf("max concurrent runs = %d, want 2 across the whole drain", got)
	}
}

func TestPerAgentTimeout(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 1, PerAgentTimeout: 50 * time.Millisecond})
	m := &mockRunner{gate: make(chan struct{})} // 永不关闭：阻塞到超时
	s.runner = m

	sub, err := s.Spawn(context.Background(), AgentSpec{Goal: "slow"})
	if err != nil {
		t.Fatalf("Spawn: %v", err)
	}
	waitFor(t, func() bool { return sub.Status() == StatusRunning }, waitTimeout)
	waitFor(t, func() bool { return sub.Status() == StatusAborted }, waitTimeout)
	if !strings.Contains(sub.Summary(), "超时") {
		t.Fatalf("Summary = %q, want a timeout notice", sub.Summary())
	}
	if !errors.Is(sub.Err(), context.DeadlineExceeded) {
		t.Fatalf("Err = %v, want DeadlineExceeded", sub.Err())
	}
}

func TestAbort(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 1})
	m := &mockRunner{gate: make(chan struct{})}
	s.runner = m

	sub, err := s.Spawn(context.Background(), AgentSpec{Goal: "x"})
	if err != nil {
		t.Fatalf("Spawn: %v", err)
	}
	waitFor(t, func() bool { return sub.Status() == StatusRunning }, waitTimeout)

	if err := s.Abort(context.Background(), sub.ID); err != nil {
		t.Fatalf("Abort: %v", err)
	}
	waitFor(t, func() bool { return sub.Status() == StatusAborted }, waitTimeout)

	// 中止一个未知 id 报告 ErrAgentNotFound。
	if err := s.Abort(context.Background(), "ghost"); !errors.Is(err, ErrAgentNotFound) {
		t.Fatalf("Abort(ghost) err = %v, want ErrAgentNotFound", err)
	}
}

func TestCancelAll(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 2})
	m := &mockRunner{gate: make(chan struct{})}
	s.runner = m

	ctx := context.Background()
	var subs []*SubAgent
	for i := 0; i < 3; i++ {
		sub, err := s.Spawn(ctx, AgentSpec{Goal: fmt.Sprintf("task %d", i)})
		if err != nil {
			t.Fatalf("Spawn %d: %v", i, err)
		}
		subs = append(subs, sub)
	}
	waitFor(t, func() bool { return countStatus(subs, StatusRunning) >= 1 }, waitTimeout)

	if err := s.CancelAll(ctx); err != nil {
		t.Fatalf("CancelAll: %v", err)
	}
	for _, sub := range subs {
		waitFor(t, func() bool { return sub.Status() == StatusAborted }, waitTimeout)
	}
}