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 is a configurable runner for scheduler tests. It counts concurrent
// invocations (for the parallel-limit test), captures the messages it receives
// (for the context-isolation test), and can block on a gate channel until closed
// or until ctx is done (to hold a sub-agent "running" for abort/timeout tests).
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 is a minimal model.Adapter so a real loop.Loop can be driven in
// tests without a provider. Chat returns a scripted final response.
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 builds a Scheduler over a fresh in-process bus, wiring the default
// comm probe. The bus is closed on test cleanup. Callers set sched.runner after
// the call, or pass Options.Loop to drive a real 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 reads events from ch until a terminal event or the deadline
// elapses, returning the event names in order.
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)
	}

	// Wait until the two parallel slots are occupied, then confirm never more
	// than MaxParallel run concurrently.
	waitFor(t, func() bool { return countStatus(subs, StatusRunning) == 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)
	}
}

func TestPerAgentTimeout(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 1, PerAgentTimeout: 50 * time.Millisecond})
	m := &mockRunner{gate: make(chan struct{})} // never closed: block until timeout
	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)

	// Aborting an unknown id reports 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)
	}
}