iop/packages/go/agenttask/reconcile.go

package agenttask

import (
	"context"
	"errors"
	"fmt"
	"maps"
	"reflect"
	"slices"
	"sort"
	"sync"
)

func (m *Manager) Reconcile(ctx context.Context) (resultErr error) {
	m.reconcileMu.Lock()
	defer m.reconcileMu.Unlock()

	ctx, deliveryErrors := withDeliveryErrors(ctx)
	defer func() {
		resultErr = errors.Join(resultErr, deliveryErrors.Err())
	}()
	if err := m.flushPendingEvents(ctx); err != nil {
		recordDeliveryError(ctx, err)
		return nil
	}

	deviceClaim, err := m.claimDevice(ctx)
	if err != nil {
		return err
	}
	ownedCtx, leases := m.maintainLeases(ctx, deviceClaim)
	defer func() {
		// Stop and join the supervisor before releasing anything. Every release
		// uses the immutable claim captured at acquisition, never a reloaded
		// current token that could belong to a successor.
		leases.Close()
		for _, claim := range leases.snapshotClaims() {
			m.releaseExact(unfencedLeaseContext(ownedCtx), claim)
		}
	}()

	if err := m.reconcileCheckpoint(ownedCtx); err != nil {
		return err
	}
	if err := deliveryErrors.Err(); err != nil {
		return nil
	}
	active, err := m.observeWorkflows(ownedCtx)
	if err != nil {
		return err
	}
	if err := deliveryErrors.Err(); err != nil {
		return nil
	}
	claimed := make([]ProjectID, 0, len(active))
	claimedWorkspaces := make(map[WorkspaceID]struct{}, len(active))
	for _, projectID := range active {
		if err := leases.Validate(ownedCtx); err != nil {
			return err
		}
		projectClaim, claimErr := m.claimProject(ownedCtx, projectID)
		if claimErr != nil {
			return claimErr
		}
		if projectClaim == nil {
			m.emit(ownedCtx, Event{
				Type:      EventBlocked,
				ProjectID: projectID,
				Detail:    string(BlockerDuplicateProjectLease),
			})
			if err := deliveryErrors.Err(); err != nil {
				return nil
			}
			continue
		}
		leases.Add(projectClaim)
		releaseProjectClaim := func() {
			leases.Remove(projectClaim)
			m.releaseExact(ownedCtx, projectClaim)
		}
		state, loadErr := m.load(ownedCtx)
		if loadErr != nil {
			releaseProjectClaim()
			return loadErr
		}
		project := state.Projects[projectID]
		if _, alreadyClaimed := claimedWorkspaces[project.WorkspaceID]; !alreadyClaimed {
			workspaceClaim, workspaceErr := m.claimWorkspace(ownedCtx, project.WorkspaceID)
			if workspaceErr != nil {
				releaseProjectClaim()
				return workspaceErr
			}
			if workspaceClaim == nil {
				m.blockProject(ownedCtx, projectID, Blocker{
					Code:    BlockerDuplicateWorkspaceCall,
					Message: "workspace invocation lease is held by another live owner",
				})
				releaseProjectClaim()
				continue
			}
			leases.Add(workspaceClaim)
			claimedWorkspaces[project.WorkspaceID] = struct{}{}
		}
		claimed = append(claimed, projectID)
	}
	if len(claimed) == 0 {
		return nil
	}
	projectContexts := make(map[ProjectID]context.Context, len(claimed))
	projectCleanups := make([]func(), 0, len(claimed))
	for _, projectID := range claimed {
		projectCtx, cleanup := m.beginProjectRun(ownedCtx, projectID)
		projectContexts[projectID] = projectCtx
		projectCleanups = append(projectCleanups, cleanup)
	}
	defer func() {
		for _, cleanup := range projectCleanups {
			cleanup()
		}
	}()
	var reconcileErrors []error
	for round := 0; round < 10_000; round++ {
		if err := leases.Validate(ownedCtx); err != nil {
			return errors.Join(append(reconcileErrors, err)...)
		}
		if err := m.refreshDependencies(ownedCtx, claimed); err != nil {
			return errors.Join(append(reconcileErrors, err)...)
		}
		if err := deliveryErrors.Err(); err != nil {
			return errors.Join(reconcileErrors...)
		}
		candidates, err := m.runnableWorks(ownedCtx, claimed)
		if err != nil {
			return errors.Join(append(reconcileErrors, err)...)
		}
		if len(candidates) > 0 {
			var wait sync.WaitGroup
			errs := make(chan error, len(candidates))
			for _, candidate := range candidates {
				candidate := candidate
				wait.Add(1)
				go func() {
					defer wait.Done()
					projectCtx := projectContexts[candidate.ProjectID]
					if runErr := m.runWork(projectCtx, candidate.ProjectID, candidate.WorkUnitID); runErr != nil {
						errs <- runErr
					}
				}()
			}
			wait.Wait()
			close(errs)
			for runErr := range errs {
				reconcileErrors = append(reconcileErrors, runErr)
			}
			if err := deliveryErrors.Err(); err != nil {
				return errors.Join(reconcileErrors...)
			}
		}
		integrated, integrationErr := m.integratePending(ownedCtx, claimed, leases)
		if integrationErr != nil {
			reconcileErrors = append(reconcileErrors, integrationErr)
		}
		if err := deliveryErrors.Err(); err != nil {
			return errors.Join(reconcileErrors...)
		}
		if len(candidates) == 0 && !integrated {
			break
		}
		if ownedCtx.Err() != nil {
			return errors.Join(append(reconcileErrors, ownedCtx.Err())...)
		}
	}
	if err := m.refreshProjectStatuses(ownedCtx, claimed); err != nil {
		return errors.Join(append(reconcileErrors, err)...)
	}
	if err := deliveryErrors.Err(); err != nil {
		return errors.Join(reconcileErrors...)
	}
	return errors.Join(reconcileErrors...)
}

func (m *Manager) reconcileCheckpoint(ctx context.Context) error {
	state, err := m.load(ctx)
	if err != nil {
		return err
	}
	projectIDs := make([]ProjectID, 0, len(state.Projects))
	for projectID := range state.Projects {
		projectIDs = append(projectIDs, projectID)
	}
	sort.Slice(projectIDs, func(left, right int) bool {
		return projectIDs[left] < projectIDs[right]
	})
	for _, projectID := range projectIDs {
		project := state.Projects[projectID]
		if project.Status != ProjectStatusRunning {
			continue
		}
		workIDs := make([]WorkUnitID, 0, len(project.Works))
		for workID := range project.Works {
			workIDs = append(workIDs, workID)
		}
		sort.Slice(workIDs, func(left, right int) bool {
			return workIDs[left] < workIDs[right]
		})
		for _, workID := range workIDs {
			work := project.Works[workID]
			needsExecution := work.State == WorkStateDispatching
			_, hasCompletion := work.Locators[LocatorCompletion]
			needsCompletion := work.State == WorkStateCompleted && hasCompletion
			if !needsExecution && !needsCompletion {
				continue
			}
			observation, inspectErr := m.recovery.Inspect(ctx, RecoveryRequest{
				Project: project, Work: work, Locators: maps.Clone(work.Locators),
			})
			if inspectErr != nil {
				m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
					Code:    BlockerAmbiguousCheckpoint,
					Message: "durable locator inspection failed: " + inspectErr.Error(),
				})
				continue
			}
			if err := validateRecoveryObservation(project, work, observation); err != nil {
				m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
					Code: BlockerStaleCheckpoint, Message: err.Error(),
				})
				continue
			}
			if needsExecution {
				if err := m.applyExecutionRecovery(ctx, project, work, observation); err != nil {
					return err
				}
				continue
			}
			switch observation.Completion {
			case RecoveryCompletionComplete:
				if err := m.changeWork(ctx, projectID, workID, func(current *WorkRecord) error {
					if !sameRecoveryGeneration(*current, work) {
						return checkpointError(
							BlockerStaleCheckpoint,
							fmt.Sprintf("work %q changed during completion recovery", workID),
						)
					}
					current.CompletionVerified = true
					current.Blocker = nil
					resetFailure(current, FailureStageRecovery)
					return nil
				}); err != nil {
					return err
				}
			case RecoveryCompletionPartial:
				m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
					Code:    BlockerPartialCompletion,
					Message: "completion archive is partial and requires exact host reconciliation",
				})
			case RecoveryCompletionUnknown, RecoveryCompletionAmbiguous:
				m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
					Code:    BlockerAmbiguousCheckpoint,
					Message: "completion locator cannot prove a complete archive",
				})
			}
		}
	}
	return nil
}

func cloneRecoveredSubmission(src *Submission) Submission {
	if src == nil {
		return Submission{}
	}
	sub := *src
	sub.Metadata = maps.Clone(src.Metadata)
	sub.Locators = slices.Clone(src.Locators)
	return sub
}

func (m *Manager) applyExecutionRecovery(
	ctx context.Context,
	project ProjectRecord,
	work WorkRecord,
	observation RecoveryObservation,
) error {
	projectID := project.ProjectID
	workID := work.Unit.ID
	switch observation.Execution {
	case RecoveryExecutionLive:
		return nil
	case RecoveryExecutionSubmitted:
		if observation.Submission == nil {
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code:    BlockerAmbiguousCheckpoint,
				Message: "finished execution has no exact submission identity",
			})
			return nil
		}
		submission := cloneRecoveredSubmission(observation.Submission)
		if err := validateSubmission(project, work, submission); err != nil {
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code: BlockerStaleCheckpoint, Message: err.Error(),
			})
			return nil
		}
		if !submission.Ready {
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code:    BlockerSubmissionIncomplete,
				Message: "worker submission did not pass the provider-neutral completeness gate",
			})
			return nil
		}
		if blocker := m.gateSubmissionEvidence(ctx, project, work, submission); blocker != nil {
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, *blocker)
			return nil
		}
		return m.changeWork(ctx, projectID, workID, func(current *WorkRecord) error {
			if !sameRecoveryGeneration(*current, work) {
				return checkpointError(
					BlockerStaleCheckpoint,
					fmt.Sprintf("work %q changed during execution recovery", workID),
				)
			}
			if err := transitionWork(current, WorkStateSubmitted); err != nil {
				return err
			}
			if err := transitionWork(current, WorkStateReviewing); err != nil {
				return err
			}
			current.Submission = &submission
			for _, locator := range submission.Locators {
				current.Locators[locator.Kind] = locator
			}
			resetFailure(current, FailureStageRecovery)
			return nil
		})
	case RecoveryExecutionAbsent:
		if _, hasProcess := work.Locators[LocatorProcess]; hasProcess {
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code:    BlockerAmbiguousCheckpoint,
				Message: "persisted process locator disappeared without terminal evidence",
			})
			return nil
		}
		if _, hasSession := work.Locators[LocatorSession]; hasSession {
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code:    BlockerAmbiguousCheckpoint,
				Message: "persisted session locator disappeared without terminal evidence",
			})
			return nil
		}
		return m.changeWork(ctx, projectID, workID, func(current *WorkRecord) error {
			if !sameRecoveryGeneration(*current, work) {
				return checkpointError(
					BlockerStaleCheckpoint,
					fmt.Sprintf("work %q changed during absent-execution recovery", workID),
				)
			}
			return transitionWork(current, WorkStateReady)
		})
	case RecoveryExecutionExited:
		m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
			Code:    BlockerAmbiguousCheckpoint,
			Message: "provider exited without a durable submission result",
		})
	case RecoveryExecutionAmbiguous:
		m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
			Code:    BlockerAmbiguousCheckpoint,
			Message: "process or session locator is ambiguous",
		})
	default:
		m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
			Code:    BlockerCorruptCheckpoint,
			Message: fmt.Sprintf("unsupported execution recovery state %q", observation.Execution),
		})
	}
	return nil
}

func validateRecoveryObservation(
	project ProjectRecord,
	work WorkRecord,
	observation RecoveryObservation,
) error {
	if observation.ProjectID != project.ProjectID ||
		observation.WorkspaceID != project.WorkspaceID ||
		observation.WorkUnitID != work.Unit.ID ||
		observation.AttemptID != work.AttemptID {
		return fmt.Errorf("agenttask: recovery observation durable identity mismatch")
	}
	return nil
}

func sameRecoveryGeneration(current, observed WorkRecord) bool {
	return current.Unit.ID == observed.Unit.ID &&
		current.Attempt == observed.Attempt &&
		current.AttemptID == observed.AttemptID &&
		current.State == observed.State &&
		reflect.DeepEqual(current.Locators, observed.Locators)
}

type runnableWork struct {
	ProjectID  ProjectID
	WorkUnitID WorkUnitID
	Ordinal    DispatchOrdinal
}

func (m *Manager) refreshDependencies(ctx context.Context, active []ProjectID) error {
	activeSet := make(map[ProjectID]struct{}, len(active))
	for _, id := range active {
		activeSet[id] = struct{}{}
	}
	events, err := mutateDecision(m, ctx, func(state *ManagerState) ([]Event, error) {
		var events []Event
		projectIDs := make([]ProjectID, 0, len(activeSet))
		for id := range activeSet {
			projectIDs = append(projectIDs, id)
		}
		sort.Slice(projectIDs, func(left, right int) bool {
			return projectIDs[left] < projectIDs[right]
		})
		for _, projectID := range projectIDs {
			project := state.Projects[projectID]
			if project.Status != ProjectStatusRunning || project.Workflow == nil {
				continue
			}
			workIDs := make([]WorkUnitID, 0, len(project.Works))
			for id := range project.Works {
				workIDs = append(workIDs, id)
			}
			sort.Slice(workIDs, func(left, right int) bool {
				return workIDs[left] < workIDs[right]
			})
			for _, workID := range workIDs {
				work := project.Works[workID]
				if work.State != WorkStateObserved {
					continue
				}
				dependency := evaluateDependencies(work.Unit, *project.Workflow, project.Works)
				switch dependency.Status {
				case dependencyReady:
					if err := transitionWork(&work, WorkStateReady); err != nil {
						return nil, err
					}
					if work.DispatchOrdinal == 0 {
						state.NextOrdinal++
						work.DispatchOrdinal = state.NextOrdinal
					}
					if work.Attempt == 0 {
						work.Attempt = 1
						work.AttemptID = attemptID(work.Unit.ID, work.Attempt)
					}
					work.Blocker = nil
					work.UpdatedAt = m.clock.Now()
					events = append(events, Event{
						Type:             EventDependencyReady,
						ProjectID:        projectID,
						WorkspaceID:      project.WorkspaceID,
						WorkUnitID:       workID,
						CommandID:        project.Intent.CommandID,
						WorkflowRevision: project.Intent.WorkflowRevision,
						AttemptID:        work.AttemptID,
						Ordinal:          work.DispatchOrdinal,
						State:            work.State,
						WriteSetKind:     work.Unit.WriteSetKind,
						IsolationMode:    work.Unit.IsolationMode,
					})
				case dependencyMissing:
					blockWorkDependency(&work, BlockerDependencyMissing, dependency.Ref)
				case dependencyAmbiguous:
					blockWorkDependency(&work, BlockerDependencyAmbiguous, dependency.Ref)
				case dependencyBlocked:
					blockWorkDependency(&work, BlockerDependencyBlocked, dependency.Ref)
				case dependencyWaiting:
					continue
				}
				project.Works[workID] = work
			}
			state.Projects[projectID] = project
		}
		return events, nil
	})
	if err != nil {
		return err
	}
	for _, event := range events {
		m.emit(ctx, event)
	}
	return nil
}

func blockWorkDependency(work *WorkRecord, code BlockerCode, ref string) {
	_ = transitionWork(work, WorkStateBlocked)
	work.Blocker = &Blocker{
		Code:    code,
		Message: fmt.Sprintf("explicit predecessor %q is %s", ref, code),
	}
}

func (m *Manager) runnableWorks(
	ctx context.Context,
	active []ProjectID,
) ([]runnableWork, error) {
	state, err := m.load(ctx)
	if err != nil {
		return nil, err
	}
	var candidates []runnableWork
	for _, projectID := range active {
		project := state.Projects[projectID]
		if project.Status != ProjectStatusRunning {
			continue
		}
		for workID, work := range project.Works {
			if work.State != WorkStateReady && work.State != WorkStateReviewing {
				continue
			}
			candidates = append(candidates, runnableWork{
				ProjectID: projectID, WorkUnitID: workID, Ordinal: work.DispatchOrdinal,
			})
		}
	}
	sort.Slice(candidates, func(left, right int) bool {
		if candidates[left].Ordinal != candidates[right].Ordinal {
			return candidates[left].Ordinal < candidates[right].Ordinal
		}
		if candidates[left].ProjectID != candidates[right].ProjectID {
			return candidates[left].ProjectID < candidates[right].ProjectID
		}
		return candidates[left].WorkUnitID < candidates[right].WorkUnitID
	})
	return candidates, nil
}

func (m *Manager) refreshProjectStatuses(ctx context.Context, active []ProjectID) error {
	events, err := mutateDecision(m, ctx, func(state *ManagerState) ([]Event, error) {
		var events []Event
		for _, projectID := range active {
			project := state.Projects[projectID]
			if project.Status == ProjectStatusStopped {
				continue
			}
			selected := 0
			completed := 0
			activeWork := 0
			for _, work := range project.Works {
				selected++
				switch {
				case work.State == WorkStateCompleted && work.CompletionVerified:
					completed++
				case !work.State.Terminal():
					activeWork++
				}
			}
			switch {
			case selected > 0 && completed == selected:
				project.Status = ProjectStatusCompleted
				project.Blocker = nil
			case activeWork > 0:
				project.Status = ProjectStatusRunning
			default:
				project.Status = ProjectStatusBlocked
			}
			project.UpdatedAt = m.clock.Now()
			state.Projects[projectID] = project
			if project.Status == ProjectStatusCompleted {
				var cmdID CommandID
				var wfRev WorkflowRevision
				if project.Intent != nil {
					cmdID = project.Intent.CommandID
					wfRev = project.Intent.WorkflowRevision
				}
				events = append(events, Event{
					Type:             EventCompleted,
					ProjectID:        projectID,
					WorkspaceID:      project.WorkspaceID,
					CommandID:        cmdID,
					WorkflowRevision: wfRev,
				})
			}
		}
		return events, nil
	})
	if err != nil {
		return err
	}
	for _, event := range events {
		m.emit(ctx, event)
	}
	return nil
}