Cyrene_v2/internal/runtime/orchestration/scheduler.go

package orchestration

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/orchestration/audit"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/orchestration/comm"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
)

// runner executes one sub-agent's loop. A *loop.Loop satisfies it; tests inject a
// mock so the scheduler can be exercised without a model backend.
type runner interface {
	Run(ctx context.Context, sessionID string, messages []model.Message) (*loop.Result, error)
}

// Scheduler creates and runs sub-agents in the background, bounding concurrency,
// applying a per-agent timeout and publishing progress over the agent-to-agent
// bus. It is safe for concurrent use.
//
// The lifecycle is: Spawn registers a sub-agent and launches it asynchronously;
// the sub-agent's goroutine acquires a parallel slot (gated by MaxParallel),
// transitions to running, runs the injected loop template in its isolated
// context under the per-agent timeout, writes a summary and publishes a terminal
// progress event. Abort/CancelAll interrupt running or queued sub-agents. Join
// blocks until every sub-agent is terminal and returns the merged outcome.
type Scheduler struct {
	maxParallel     int
	perAgentTimeout time.Duration
	runner          runner
	bus             bus.Bus
	comm            *comm.Probe
	memory          *memory.Store
	tracer          *trace.Tracer
	approver        *audit.Approver
	log             *slog.Logger
	agentID         string

	mu      sync.Mutex
	agents  map[string]*SubAgent
	active  int
	allDone chan struct{}
	allOne  sync.Once
	sem     chan struct{}
	seq     atomic.Uint64
	prefix  string
}

// noopRunner is the fallback runner used when no *loop.Loop template was
// configured. It makes every sub-agent fail fast with ErrNoRunner rather than
// silently doing nothing.
type noopRunner struct{}

// Run returns a wrapped ErrNoRunner.
func (noopRunner) Run(context.Context, string, []model.Message) (*loop.Result, error) {
	return nil, fmt.Errorf("orchestration: run sub-agent: %w", ErrNoRunner)
}

// New builds a Scheduler from opts, applying defaults. MaxParallel non-positive
// selects DefaultMaxParallel; a nil Loop is replaced by a no-op runner; a nil
// Comm is built from Bus (when present) using AgentID (default "orchestrator"); a
// nil Logger is replaced by a discard logger tagged with source "orchestration".
func New(opts Options) *Scheduler {
	maxParallel := opts.MaxParallel
	if maxParallel <= 0 {
		maxParallel = DefaultMaxParallel
	}
	agentID := opts.AgentID
	if agentID == "" {
		agentID = defaultAgentID
	}

	s := &Scheduler{
		maxParallel:     maxParallel,
		perAgentTimeout: opts.PerAgentTimeout,
		runner:          wrapLoop(opts.Loop),
		bus:             opts.Bus,
		memory:          opts.Memory,
		tracer:          opts.Tracer,
		approver:        opts.Approver,
		log:             sourcedLog(opts.Logger),
		agentID:         agentID,
		agents:          make(map[string]*SubAgent),
		allDone:         make(chan struct{}),
		sem:             make(chan struct{}, maxParallel),
		prefix:          randPrefix(),
	}
	if opts.Comm != nil {
		s.comm = opts.Comm
	} else if opts.Bus != nil {
		s.comm = comm.New(opts.Bus, agentID)
	}
	return s
}

// Bus returns the injected event bus (may be nil).
func (s *Scheduler) Bus() bus.Bus { return s.bus }

// Comm returns the injected (or auto-built) comm probe (may be nil).
func (s *Scheduler) Comm() *comm.Probe { return s.comm }

// Spawn creates a sub-agent from spec, registers it, publishes a "dispatched"
// progress event and launches its execution in the background. It returns
// immediately so the front session is not blocked; the sub-agent runs
// asynchronously and reports progress over the bus. It returns ErrMissingGoal
// when spec.Goal is blank.
//
// The sub-agent's Context is a deep copy of spec.Context and its WorkingMemory is
// a fresh, isolated window, so it never shares state with the front session.
func (s *Scheduler) Spawn(ctx context.Context, spec AgentSpec) (*SubAgent, error) {
	if strings.TrimSpace(spec.Goal) == "" {
		return nil, fmt.Errorf("orchestration: spawn: %w", ErrMissingGoal)
	}
	if ctx == nil {
		ctx = context.Background()
	}

	sub := &SubAgent{
		ID:            s.nextID(),
		Goal:          spec.Goal,
		RoleID:        spec.RoleID,
		Mode:          spec.Mode,
		AttachInfo:    spec.AttachInfo,
		Data:          spec.Data,
		Context:       cloneMessages(spec.Context),
		WorkingMemory: memory.NewWorkingMemory(0),
		status:        StatusPending,
		done:          make(chan struct{}),
	}

	s.mu.Lock()
	if s.active == 0 {
		s.allDone = make(chan struct{})
		s.allOne = sync.Once{}
	}
	s.active++
	s.agents[sub.ID] = sub
	s.mu.Unlock()

	s.publish(EventDispatched, sub.progress())
	s.launch(ctx, sub)
	return sub, nil
}

// Run explicitly starts every registered sub-agent that has not started yet. It
// is idempotent and is provided for callers that prefer to defer execution; Spawn
// already auto-launches, so this is normally a no-op on freshly spawned batches.
func (s *Scheduler) Run(ctx context.Context) {
	if ctx == nil {
		ctx = context.Background()
	}
	s.mu.Lock()
	subs := make([]*SubAgent, 0, len(s.agents))
	for _, sub := range s.agents {
		subs = append(subs, sub)
	}
	s.mu.Unlock()
	for _, sub := range subs {
		s.launch(ctx, sub)
	}
}

// Join blocks until every spawned sub-agent is terminal (or ctx is canceled) and
// returns the merged outcome. On ctx cancellation it returns the partial merged
// outcome together with the wrapped ctx error.
func (s *Scheduler) Join(ctx context.Context) (*MergedResult, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	s.mu.Lock()
	ch := s.allDone
	s.mu.Unlock()

	select {
	case <-ch:
		return s.merge(), nil
	case <-ctx.Done():
		return s.merge(), ctx.Err()
	}
}

// Abort cancels the sub-agent with the given id. If it is running its loop is
// interrupted; if it is queued it is released without running; if it already
// finished the call is a no-op. It returns ErrAgentNotFound for an unknown id.
func (s *Scheduler) Abort(ctx context.Context, id string) error {
	if ctx == nil {
		ctx = context.Background()
	}
	s.mu.Lock()
	sub, ok := s.agents[id]
	s.mu.Unlock()
	if !ok {
		return fmt.Errorf("orchestration: abort: %w", ErrAgentNotFound)
	}
	sub.cancelRun()
	return nil
}

// CancelAll cancels every spawned sub-agent, queued or running. It is the
// anti-runaway escape hatch: after it returns each sub-agent converges to the
// aborted state.
func (s *Scheduler) CancelAll(ctx context.Context) error {
	if ctx == nil {
		ctx = context.Background()
	}
	s.mu.Lock()
	subs := make([]*SubAgent, 0, len(s.agents))
	for _, sub := range s.agents {
		subs = append(subs, sub)
	}
	s.mu.Unlock()
	for _, sub := range subs {
		sub.cancelRun()
	}
	return nil
}

// Agents returns a snapshot of every registered sub-agent in spawn order. It is
// safe to read the returned snapshots at any time.
func (s *Scheduler) Agents() []SubAgentResult {
	s.mu.Lock()
	subs := make([]*SubAgent, 0, len(s.agents))
	for _, sub := range s.agents {
		subs = append(subs, sub)
	}
	s.mu.Unlock()

	out := make([]SubAgentResult, 0, len(subs))
	for _, sub := range subs {
		out = append(out, sub.Snapshot())
	}
	return out
}

// Approve triggers a main-agent review-and-approve (代审) of a sub-agent's
// sensitive action through the injected *audit.Approver. Whether an action is
// sensitive and whether it is granted is decided by the upper policy layer; this
// method only records the audit entry. It returns ErrNoApprover when no approver
// was injected.
func (s *Scheduler) Approve(ctx context.Context, actor, action, describe string) (*audit.Approval, error) {
	if s.approver == nil {
		return nil, fmt.Errorf("orchestration: approve: %w", ErrNoApprover)
	}
	return s.approver.ReviewAndApprove(ctx, actor, action, describe)
}

// SinkMemory persists a task memory item for a sub-agent through the injected
// *memory.Store. The sub-agent's RoleID is used as the persona/role key; content
// and the scope/kind classification (user/persona/project) are supplied by the
// caller, keeping scope policy at the call site. It returns ErrNoMemory when no
// store was injected.
func (s *Scheduler) SinkMemory(ctx context.Context, sub *SubAgent, content string, scope memory.Scope, kind memory.Kind) (memory.MemoryItem, error) {
	if s.memory == nil {
		return memory.MemoryItem{}, fmt.Errorf("orchestration: sink memory: %w", ErrNoMemory)
	}
	if sub == nil {
		return memory.MemoryItem{}, fmt.Errorf("orchestration: sink memory: %w", ErrAgentNotFound)
	}
	item := memory.MemoryItem{
		PersonaID: sub.RoleID,
		RoleID:    sub.RoleID,
		Scope:     scope,
		Kind:      kind,
		Content:   content,
	}
	return s.memory.Add(ctx, item)
}

// launch starts the sub-agent's background execution exactly once. It is called
// by Spawn (auto-launch) and by Run (explicit start); launchOnce guarantees a
// single run even if both happen.
func (s *Scheduler) launch(ctx context.Context, sub *SubAgent) {
	sub.launchOnce.Do(func() {
		go s.run(ctx, sub)
	})
}

// run drives one sub-agent: acquire a parallel slot, transition to running, run
// the loop in the sub-agent's isolated context under the per-agent timeout, then
// finalize with a summary and a terminal progress event. It always reaches
// exactly one finish, so active is decremented once and Join unblocks exactly
// when every sub-agent is terminal.
func (s *Scheduler) run(ctx context.Context, sub *SubAgent) {
	if ctx == nil {
		ctx = context.Background()
	}

	select {
	case s.sem <- struct{}{}:
		defer func() { <-s.sem }()
	case <-ctx.Done():
		s.finish(sub, StatusAborted, canceledSummary(), ctx.Err())
		return
	case <-sub.done:
		s.finish(sub, StatusAborted, canceledSummary(), context.Canceled)
		return
	}

	if sub.aborted() || !sub.transitionRunning() {
		s.finish(sub, StatusAborted, canceledSummary(), context.Canceled)
		return
	}
	s.publish(EventStarted, sub.progress())

	runCtx, cancel := s.subCtx(ctx)
	runCtx, span := s.subSpan(runCtx, sub)
	sub.setCancel(cancel)
	defer cancel()

	res, err := s.runner.Run(runCtx, sub.ID, sub.Context)
	if err != nil {
		s.spanErr(span, err)
		s.abortOrFail(sub, err)
		return
	}
	s.spanOK(span, res)
	sub.setResult(res)
	s.finish(sub, StatusDone, doneSummary(res), nil)
}

// abortOrFail classifies an error after a sub-agent's run: a deadline exceeded or
// a cancellation yields the aborted status (timeout/stop), anything else is a
// real failure.
func (s *Scheduler) abortOrFail(sub *SubAgent, err error) {
	switch {
	case errors.Is(err, context.DeadlineExceeded):
		s.finish(sub, StatusAborted, timeoutSummary(s.perAgentTimeout), err)
	case errors.Is(err, context.Canceled):
		s.finish(sub, StatusAborted, canceledSummary(), err)
	default:
		s.finish(sub, StatusFailed, failedSummary(err), err)
	}
}

// finish finalizes a sub-agent exactly once: it records the terminal status and
// summary, publishes the terminal progress event and decrements the active count,
// closing the allDone channel when the last sub-agent finalizes.
func (s *Scheduler) finish(sub *SubAgent, status Status, summary string, err error) {
	sub.finishOnce.Do(func() {
		sub.setOutcome(status, summary, err)
		s.publish(eventFor(status), sub.progress())

		s.mu.Lock()
		s.active--
		if s.active == 0 {
			s.allOne.Do(func() { close(s.allDone) })
		}
		s.mu.Unlock()
	})
}

// subCtx derives the run context, applying the per-agent timeout when configured.
func (s *Scheduler) subCtx(parent context.Context) (context.Context, context.CancelFunc) {
	if s.perAgentTimeout > 0 {
		return context.WithTimeout(parent, s.perAgentTimeout)
	}
	return context.WithCancel(parent)
}

// subSpan starts a sub-agent span when a tracer is configured, chaining it as a
// child of any span already in ctx. It returns the ctx carrying the span and the
// span itself (nil when the tracer is nil).
func (s *Scheduler) subSpan(ctx context.Context, sub *SubAgent) (context.Context, *trace.Span) {
	if s.tracer == nil {
		return ctx, nil
	}
	return s.tracer.StartSpan(ctx, "subagent."+sub.ID, trace.SpanKindSubagent)
}

// spanOK finishes a sub-agent span as ok, recording the result text.
func (s *Scheduler) spanOK(span *trace.Span, res *loop.Result) {
	if span == nil {
		return
	}
	if res != nil {
		_ = span.SetProp("text", res.Text)
	}
	span.Finish(trace.SpanStatusOK)
}

// spanErr finishes a sub-agent span as an error, recording the cause.
func (s *Scheduler) spanErr(span *trace.Span, err error) {
	if span == nil {
		return
	}
	_ = span.SetProp("error", err.Error())
	span.Finish(trace.SpanStatusError)
}

// publish broadcasts a progress event over the bus. It is a no-op when no comm
// probe is configured. Publishing uses a background context so a canceled run
// still delivers its terminal event before the sub-agent is done.
func (s *Scheduler) publish(name string, data Progress) {
	if s.comm == nil {
		return
	}
	if err := s.comm.PublishEvent(context.Background(), comm.Event{Name: name, Data: data}); err != nil {
		s.log.Warn("publish progress", "event", name, "agent_id", data.AgentID, "error", err)
	}
}

// nextID mints a unique sub-agent id from the scheduler's random prefix and an
// atomic sequence. It is safe for concurrent use.
func (s *Scheduler) nextID() string {
	return fmt.Sprintf("%s-%d", s.prefix, s.seq.Add(1))
}

// merge aggregates every registered sub-agent into a MergedResult. It is called by
// Join once all sub-agents are terminal.
func (s *Scheduler) merge() *MergedResult {
	s.mu.Lock()
	subs := make([]*SubAgent, 0, len(s.agents))
	for _, sub := range s.agents {
		subs = append(subs, sub)
	}
	s.mu.Unlock()

	res := &MergedResult{Total: len(subs), Agents: make([]SubAgentResult, 0, len(subs))}
	for _, sub := range subs {
		snap := sub.Snapshot()
		res.Agents = append(res.Agents, snap)
		switch snap.Status {
		case StatusDone:
			res.Succeeded++
		case StatusFailed:
			res.Failed++
		case StatusAborted:
			res.Aborted++
		}
	}
	res.Summary = summarizeMerged(res)
	return res
}

// eventFor maps a terminal status to its progress event name.
func eventFor(status Status) string {
	switch status {
	case StatusDone:
		return EventDone
	case StatusFailed:
		return EventFailed
	case StatusAborted:
		return EventAborted
	default:
		return EventStarted
	}
}

// wrapLoop adapts a *loop.Loop into the runner interface, substituting a no-op
// runner (fail-fast) when the template is nil.
func wrapLoop(l *loop.Loop) runner {
	if l == nil {
		return noopRunner{}
	}
	return l
}

// sourcedLog returns the base logger tagged with source "orchestration",
// replacing a nil logger with a discard one so the scheduler can always log.
func sourcedLog(l *slog.Logger) *slog.Logger {
	if l == nil {
		l = slog.New(slog.NewTextHandler(io.Discard, nil))
	}
	return logger.WithSource(l, "orchestration")
}

// prefixBytes is the number of random bytes used to build the sub-agent id
// prefix.
const prefixBytes = 4

// randPrefix returns a short random prefix for sub-agent ids, falling back to
// "sub" if the system entropy source is unavailable.
func randPrefix() string {
	var b [prefixBytes]byte
	if _, err := rand.Read(b[:]); err != nil {
		return "sub"
	}
	return "sub-" + hex.EncodeToString(b[:])
}