Cyrene_v2/internal/runtime/task/manager.go

package task

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/exec"
)

// autoBackgroundThreshold is how long a foreground task may run before Start
// promotes it to the background (DevDocs/design/04-运行时.md §8.5: "前台运行，
// 超过 2 秒自动移入后台"). It is overridable via WithAutoBackground.
const autoBackgroundThreshold = 2 * time.Second

// Sentinel errors returned by this package. Use errors.Is to test for them; they
// are returned directly or wrapped with %w so the underlying cause stays visible.
var (
	// ErrNilProvider is returned by New when no exec.Provider is supplied.
	ErrNilProvider = errors.New("task: nil exec provider")
	// ErrNilLogger is returned by New when no logger is supplied.
	ErrNilLogger = errors.New("task: nil logger")
	// ErrEmptyCommand is returned by Start when spec.Command is empty.
	ErrEmptyCommand = errors.New("task: command is empty")
	// ErrNotFound is returned by Get/Abort for an unknown task id.
	ErrNotFound = errors.New("task: not found")
	// ErrAlreadyDone is returned by Abort for a task that already reached a
	// terminal state.
	ErrAlreadyDone = errors.New("task: already finished")
	// ErrManagerClosed is returned by Start on a Manager that has been closed.
	ErrManagerClosed = errors.New("task: manager closed")
)

// DoneEvent is the payload published to the "task.done" subject when a task
// reaches a terminal state. It carries everything the agent needs to consume the
// result without a second lookup.
type DoneEvent struct {
	ID         string    `json:"id"`
	Kind       Kind      `json:"kind"`
	Status     Status    `json:"status"`
	Output     Output    `json:"output"`
	ExitCode   int       `json:"exitCode"`
	Truncated  bool      `json:"truncated"`
	Error      string    `json:"error,omitempty"`
	StartedAt  time.Time `json:"startedAt"`
	FinishedAt time.Time `json:"finishedAt"`
}

// OutputEvent is the payload published to the "task.output" subject as output
// becomes available.
type OutputEvent struct {
	ID        string `json:"id"`
	Kind      Kind   `json:"kind"`
	Status    Status `json:"status"`
	Output    Output `json:"output"`
	Truncated bool   `json:"truncated"`
}

// Manager owns the task registry and the execution lifecycle. It owns an
// in-process bus (or one injected via WithBus) and publishes task events there.
type Manager struct {
	cfg       *config.Config
	prov      exec.Provider
	log       *slog.Logger
	bus       bus.Bus
	threshold time.Duration

	ctx    context.Context
	cancel context.CancelFunc
	closed atomic.Bool

	mu     sync.RWMutex
	tasks  map[string]*Task
	nextID atomic.Uint64
	active atomic.Int64
}

// options holds the resolved construction options.
type options struct {
	bus       bus.Bus
	threshold time.Duration
}

// Option customizes a Manager built by New.
type Option func(*options)

// WithBus injects an existing bus instead of letting the Manager create one.
// This is handy for tests that share a bus and for wiring a global bus.
func WithBus(b bus.Bus) Option {
	return func(o *options) { o.bus = b }
}

// WithAutoBackground overrides the foreground→background promotion threshold.
// Zero uses the default of 2 seconds.
func WithAutoBackground(d time.Duration) Option {
	return func(o *options) { o.threshold = d }
}

// New builds a Manager. The primary call form is New(cfg, p, log); the variadic
// options only refine the defaults. The manager takes a child of a fresh
// background context, so tasks are independent of the context passed to Start
// (which is honored only for a foreground wait and is never used to kill a
// background task). Close releases everything.
func New(cfg *config.Config, p exec.Provider, log *slog.Logger, opts ...Option) (*Manager, error) {
	if p == nil {
		return nil, ErrNilProvider
	}
	if log == nil {
		return nil, ErrNilLogger
	}

	o := options{threshold: autoBackgroundThreshold}
	for _, opt := range opts {
		opt(&o)
	}
	if o.bus == nil {
		o.bus = bus.New()
	}

	ctx, cancel := context.WithCancel(context.Background())
	m := &Manager{
		cfg:       cfg,
		prov:      p,
		log:       logger.WithSource(log, "runtime/task"),
		bus:       o.bus,
		threshold: o.threshold,
		ctx:       ctx,
		cancel:    cancel,
		tasks:     make(map[string]*Task),
	}
	m.log.Debug("task manager ready", "auto_background_ms", m.threshold.Milliseconds())
	if cfg != nil {
		m.log.Debug("exec sandbox", "enabled", cfg.ExecSandbox.Enabled)
	}
	return m, nil
}

// Bus returns the bus the Manager publishes task events on. Agents subscribe to
// "task.output" and "task.done" to receive task results.
func (m *Manager) Bus() bus.Bus { return m.bus }

// Active returns the number of tasks whose run goroutine is currently executing.
func (m *Manager) Active() int { return int(m.active.Load()) }

// Start schedules a task. For a background task it returns immediately. For a
// foreground task it waits up to the auto-background threshold: if the task
// finishes first it returns the terminal task, otherwise it promotes the task to
// the background and returns immediately.
func (m *Manager) Start(ctx context.Context, spec TaskSpec) (*Task, error) {
	if err := ctx.Err(); err != nil {
		return nil, fmt.Errorf("task: start: %w", err)
	}
	if len(spec.Command) == 0 {
		return nil, fmt.Errorf("task: start: %w", ErrEmptyCommand)
	}
	if m.closed.Load() {
		return nil, fmt.Errorf("task: start: %w", ErrManagerClosed)
	}

	t := &Task{
		ID:        m.genID(),
		Kind:      kindFrom(spec.Background),
		Cmd:       append([]string(nil), spec.Command...),
		StartedAt: time.Now(),
		done:      make(chan struct{}),
	}
	t.ctx, t.cancel = context.WithCancel(m.ctx)
	t.m = m
	m.add(t)
	t.setRunning()

	req := m.buildRequest(spec)
	go m.run(t, req)

	if spec.Background {
		return t, nil
	}

	// Foreground: wait for completion, promotion, or a parent-context cancel.
	timer := time.NewTimer(m.threshold)
	defer timer.Stop()
	select {
	case <-t.done:
		return t, nil
	case <-timer.C:
		t.promoteBackground()
		return t, nil
	case <-ctx.Done():
		_ = t.abort()
		return nil, fmt.Errorf("task: start %q: %w", t.ID, ctx.Err())
	}
}

// List returns a snapshot of all tasks, newest first. The returned pointers are
// live; read their state through the accessor methods.
func (m *Manager) List(ctx context.Context) ([]*Task, error) {
	if err := ctx.Err(); err != nil {
		return nil, fmt.Errorf("task: list: %w", err)
	}
	m.mu.RLock()
	out := make([]*Task, 0, len(m.tasks))
	for _, t := range m.tasks {
		out = append(out, t)
	}
	m.mu.RUnlock()

	sort.Slice(out, func(i, j int) bool {
		return out[i].StartedAt.After(out[j].StartedAt)
	})
	return out, nil
}

// Get returns the task for id or ErrNotFound.
func (m *Manager) Get(ctx context.Context, id string) (*Task, error) {
	if err := ctx.Err(); err != nil {
		return nil, fmt.Errorf("task: get %q: %w", id, err)
	}
	t := m.get(id)
	if t == nil {
		return nil, fmt.Errorf("task: get %q: %w", id, ErrNotFound)
	}
	return t, nil
}

// TaskOutput returns the current output snapshot of task id for external
// consumers (the gateway's task_output_query WS route). It never waits for the
// task to finish: for a running background task it returns the output captured so
// far. It wraps ErrNotFound with %w so callers can classify with errors.Is.
func (m *Manager) TaskOutput(ctx context.Context, id string) (OutputEvent, error) {
	if err := ctx.Err(); err != nil {
		return OutputEvent{}, fmt.Errorf("task: output %q: %w", id, err)
	}
	t := m.get(id)
	if t == nil {
		return OutputEvent{}, fmt.Errorf("task: output %q: %w", id, ErrNotFound)
	}
	return OutputEvent{
		ID:        t.ID,
		Kind:      t.Kind,
		Status:    t.Status(),
		Output:    t.Output(),
		Truncated: t.Truncated(),
	}, nil
}

// Abort cancels the task for id. It returns ErrNotFound if there is no such task
// and ErrAlreadyDone if the task has already finished.
func (m *Manager) Abort(ctx context.Context, id string) error {
	if err := ctx.Err(); err != nil {
		return fmt.Errorf("task: abort %q: %w", id, err)
	}
	t := m.get(id)
	if t == nil {
		return fmt.Errorf("task: abort %q: %w", id, ErrNotFound)
	}
	return t.abort()
}

// Close cancels every running task, releases their goroutines, and closes the
// bus. It is idempotent and safe to call concurrently.
func (m *Manager) Close() error {
	if m.closed.CompareAndSwap(false, true) {
		m.cancel()
		return m.bus.Close()
	}
	return nil
}

// buildRequest maps a TaskSpec onto the exec request, applying defaults.
func (m *Manager) buildRequest(spec TaskSpec) exec.ExecRequest {
	req := exec.ExecRequest{
		Command:        spec.Command,
		WorkDir:        spec.WorkDir,
		Timeout:        spec.Timeout,
		MaxOutputBytes: spec.MaxOutputBytes,
	}
	if req.Timeout == 0 {
		req.Timeout = exec.DefaultTimeout
	}
	if req.MaxOutputBytes == 0 {
		req.MaxOutputBytes = exec.DefaultMaxOutputBytes
	}
	return req
}

func (m *Manager) add(t *Task) {
	m.mu.Lock()
	m.tasks[t.ID] = t
	m.mu.Unlock()
}

func (m *Manager) get(id string) *Task {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.tasks[id]
}

func (m *Manager) genID() string {
	n := m.nextID.Add(1)
	return fmt.Sprintf("task-%s-%d", time.Now().UTC().Format("20060102150405.000000000"), n)
}

func kindFrom(background bool) Kind {
	if background {
		return KindBackground
	}
	return KindForeground
}