iop/packages/go/agenttask/manager.go

package agenttask

import (
	"context"
	"errors"
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"sync"
	"time"
)

type ManagerConfig struct {
	OwnerID            string
	LeaseDuration      time.Duration
	MaxReworkAttempts  uint32
	StateWriteAttempts int
}

type Manager struct {
	config     ManagerConfig
	clock      Clock
	store      StateStore
	workflow   WorkflowAdapter
	selector   Selector
	isolation  IsolationBackend
	invoker    ProviderInvoker
	reviewer   Reviewer
	integrator Integrator
	events     EventSink
	scheduler  *Scheduler

	reconcileMu sync.Mutex
	activeMu    sync.Mutex
	activeRuns  map[ProjectID]context.CancelFunc
}

func NewManager(
	config ManagerConfig,
	clock Clock,
	store StateStore,
	workflow WorkflowAdapter,
	selector Selector,
	isolation IsolationBackend,
	invoker ProviderInvoker,
	reviewer Reviewer,
	integrator Integrator,
	events EventSink,
) (*Manager, error) {
	if err := validateIdentity("manager_owner", config.OwnerID); err != nil {
		return nil, err
	}
	if store == nil || workflow == nil || selector == nil || isolation == nil ||
		invoker == nil || reviewer == nil || integrator == nil {
		return nil, errors.New("agenttask: every execution port is required; unsafe fallback is disabled")
	}
	if clock == nil {
		clock = systemClock{}
	}
	if events == nil {
		events = nopEventSink{}
	}
	if config.LeaseDuration <= 0 {
		config.LeaseDuration = 30 * time.Second
	}
	if config.MaxReworkAttempts == 0 {
		config.MaxReworkAttempts = 3
	}
	if config.StateWriteAttempts <= 0 {
		config.StateWriteAttempts = 32
	}
	return &Manager{
		config:     config,
		clock:      clock,
		store:      store,
		workflow:   workflow,
		selector:   selector,
		isolation:  isolation,
		invoker:    invoker,
		reviewer:   reviewer,
		integrator: integrator,
		events:     events,
		scheduler:  NewScheduler(),
		activeRuns: make(map[ProjectID]context.CancelFunc),
	}, nil
}

func (m *Manager) StartProject(ctx context.Context, req StartRequest) error {
	if err := validateStartRequest(req); err != nil {
		return err
	}
	autoResume := true
	if req.AutoResumeInterrupted != nil {
		autoResume = *req.AutoResumeInterrupted
	}
	intent := StartIntent{
		CommandID:             req.CommandID,
		ProjectID:             req.ProjectID,
		WorkspaceID:           req.WorkspaceID,
		MilestoneID:           req.MilestoneID,
		WorkflowRevision:      req.WorkflowRevision,
		ConfigRevision:        req.ConfigRevision,
		GrantRevision:         req.GrantRevision,
		AutoResumeInterrupted: autoResume,
		StartedAt:             m.clock.Now(),
	}
	err := m.mutate(ctx, func(state *ManagerState) error {
		if previous, ok := state.Commands[req.CommandID]; ok {
			if sameCommandIntent(previous.Intent, intent) {
				return nil
			}
			return fmt.Errorf("agenttask: command %q was already used with different immutable input", req.CommandID)
		}
		project := state.Projects[req.ProjectID]
		if project.ProjectID != "" && project.WorkspaceID != "" &&
			project.WorkspaceID != req.WorkspaceID {
			return fmt.Errorf("agenttask: project %q workspace identity changed", req.ProjectID)
		}
		project.ProjectID = req.ProjectID
		project.WorkspaceID = req.WorkspaceID
		project.Status = ProjectStatusStarted
		project.Intent = &intent
		project.Blocker = nil
		project.UpdatedAt = m.clock.Now()
		if project.Works == nil {
			project.Works = make(map[WorkUnitID]WorkRecord)
		}
		state.Commands[req.CommandID] = CommandRecord{Intent: intent}
		state.Projects[req.ProjectID] = project
		return nil
	})
	if err == nil {
		m.emit(ctx, Event{
			Type:             EventManualStart,
			ProjectID:        req.ProjectID,
			WorkspaceID:      req.WorkspaceID,
			CommandID:        req.CommandID,
			WorkflowRevision: req.WorkflowRevision,
			Detail:           string(req.MilestoneID),
		})
	}
	return err
}

func (m *Manager) StopProject(ctx context.Context, projectID ProjectID) error {
	if err := validateIdentity("project", string(projectID)); err != nil {
		return err
	}
	m.activeMu.Lock()
	cancel := m.activeRuns[projectID]
	m.activeMu.Unlock()
	if cancel != nil {
		cancel()
	}
	var commandID CommandID
	var workflowRev WorkflowRevision
	err := m.mutate(ctx, func(state *ManagerState) error {
		project, ok := state.Projects[projectID]
		if !ok {
			return fmt.Errorf("agenttask: project %q is not registered in manager state", projectID)
		}
		if project.Intent != nil {
			commandID = project.Intent.CommandID
			workflowRev = project.Intent.WorkflowRevision
		}
		project.Status = ProjectStatusStopped
		project.Lease = nil
		project.UpdatedAt = m.clock.Now()
		for id, work := range project.Works {
			if !work.State.Terminal() {
				preStop := work.State
				if preStop != WorkStateStopped {
					work.ResumeStage = preStop
				}
				if err := transitionWork(&work, WorkStateStopped); err != nil {
					return err
				}
				work.UpdatedAt = m.clock.Now()
				project.Works[id] = work
			}
		}
		state.Projects[projectID] = project
		return nil
	})
	if err == nil {
		m.emit(ctx, Event{
			Type:             EventStopped,
			ProjectID:        projectID,
			CommandID:        commandID,
			WorkflowRevision: workflowRev,
		})
	}
	return err
}

func mutateDecision[T any](
	m *Manager,
	ctx context.Context,
	change func(*ManagerState) (T, error),
) (T, error) {
	var zero T
	var conflict error
	for range m.config.StateWriteAttempts {
		state, revision, err := m.store.Load(ctx)
		if err != nil {
			return zero, err
		}
		next := cloneState(state)
		if next.SchemaVersion != currentSchemaVersion {
			return zero, fmt.Errorf("agenttask: unsupported state schema %d", next.SchemaVersion)
		}
		result, err := change(&next)
		if err != nil {
			return zero, err
		}
		_, err = m.store.CompareAndSwap(ctx, revision, next)
		if errors.Is(err, ErrRevisionConflict) {
			conflict = err
			continue
		}
		if err != nil {
			return zero, err
		}
		return result, nil
	}
	return zero, fmt.Errorf("agenttask: state CAS retry budget exhausted: %w", conflict)
}

func (m *Manager) mutate(
	ctx context.Context,
	change func(*ManagerState) error,
) error {
	_, err := mutateDecision(m, ctx, func(state *ManagerState) (struct{}, error) {
		return struct{}{}, change(state)
	})
	return err
}

func (m *Manager) load(ctx context.Context) (ManagerState, error) {
	state, _, err := m.store.Load(ctx)
	if err != nil {
		return ManagerState{}, err
	}
	state = cloneState(state)
	if state.SchemaVersion != currentSchemaVersion {
		return ManagerState{}, fmt.Errorf("agenttask: unsupported state schema %d", state.SchemaVersion)
	}
	return state, nil
}

func durableIdentity(domain string, components ...string) string {
	var sb strings.Builder
	sb.WriteString(domain)
	for _, c := range components {
		sb.WriteString("/")
		sb.WriteString(strconv.Itoa(len(c)))
		sb.WriteString(":")
		sb.WriteString(c)
	}
	return sb.String()
}

func (m *Manager) emit(ctx context.Context, event Event) {
	if event.EventID == "" {
		event.EventID = durableIdentity(
			"event-v1",
			string(event.Type),
			string(event.ProjectID),
			string(event.WorkspaceID),
			string(event.WorkUnitID),
			string(event.CommandID),
			string(event.WorkflowRevision),
			string(event.AttemptID),
			strconv.FormatUint(uint64(event.Ordinal), 10),
			string(event.ChangeSetID),
			event.ChangeSetRevision,
			strconv.FormatUint(uint64(event.IntegrationAttempt), 10),
			string(event.State),
			event.Detail,
		)
	}
	event.Timestamp = m.clock.Now()
	_ = m.events.Emit(ctx, event)
}

func sameCommandIntent(left, right StartIntent) bool {
	left.StartedAt = time.Time{}
	right.StartedAt = time.Time{}
	return reflect.DeepEqual(left, right)
}

func (m *Manager) beginProjectRun(
	parent context.Context,
	projectID ProjectID,
) (context.Context, func()) {
	ctx, cancel := context.WithCancel(parent)
	m.activeMu.Lock()
	if previous := m.activeRuns[projectID]; previous != nil {
		previous()
	}
	m.activeRuns[projectID] = cancel
	m.activeMu.Unlock()
	return ctx, func() {
		cancel()
		m.activeMu.Lock()
		delete(m.activeRuns, projectID)
		m.activeMu.Unlock()
	}
}

func attemptID(workID WorkUnitID, attempt uint32) AttemptID {
	return AttemptID(durableIdentity("attempt-v1", string(workID), strconv.FormatUint(uint64(attempt), 10)))
}

func dispatchKey(projectID ProjectID, workID WorkUnitID, attempt uint32) string {
	return durableIdentity("dispatch-v1", string(projectID), string(workID), strconv.FormatUint(uint64(attempt), 10))
}

func reviewKey(projectID ProjectID, workID WorkUnitID, attempt uint32, artifact ArtifactID) string {
	return durableIdentity("review-v1", string(projectID), string(workID), strconv.FormatUint(uint64(attempt), 10), string(artifact))
}

func integrationKey(
	projectID ProjectID,
	workID WorkUnitID,
	changeSet ChangeSetIdentity,
	attempt IntegrationAttempt,
) string {
	return durableIdentity(
		"integrate-v1",
		string(projectID), string(workID), string(changeSet.ID), string(changeSet.Revision), strconv.FormatUint(uint64(attempt), 10),
	)
}