iop/packages/go/agenttask/state_machine.go

package agenttask

import (
	"errors"
	"fmt"
	"maps"
	"slices"
	"strings"

	"iop/packages/go/agentpolicy"
)

const currentSchemaVersion uint32 = StateSchemaVersion

var ErrCheckpointInvalid = errors.New("agenttask invalid durable checkpoint")

type CheckpointError struct {
	Code    BlockerCode
	Message string
}

func (e *CheckpointError) Error() string {
	return fmt.Sprintf("agenttask: %s: %s", e.Code, e.Message)
}

func (e *CheckpointError) Unwrap() error { return ErrCheckpointInvalid }

var legalWorkTransitions = map[WorkState]map[WorkState]struct{}{
	WorkStateObserved: {
		WorkStateReady: {}, WorkStateCompleted: {}, WorkStateBlocked: {}, WorkStateStopped: {},
	},
	WorkStateReady: {
		WorkStatePreparing: {}, WorkStateBlocked: {}, WorkStateStopped: {},
	},
	WorkStatePreparing: {
		WorkStateReady: {}, WorkStateDispatching: {}, WorkStateBlocked: {}, WorkStateStopped: {},
	},
	WorkStateDispatching: {
		WorkStateReady: {}, WorkStateSubmitted: {}, WorkStateBlocked: {}, WorkStateStopped: {},
	},
	WorkStateSubmitted: {
		WorkStateReviewing: {}, WorkStateBlocked: {}, WorkStateStopped: {},
	},
	WorkStateReviewing: {
		WorkStateReady: {}, WorkStatePendingIntegration: {}, WorkStateTerminalDeferred: {},
		WorkStateBlocked: {}, WorkStateStopped: {},
	},
	WorkStatePendingIntegration: {
		WorkStateIntegrating: {}, WorkStateTerminalDeferred: {}, WorkStateBlocked: {}, WorkStateStopped: {},
	},
	WorkStateIntegrating: {
		WorkStatePendingIntegration: {}, WorkStateCompleted: {}, WorkStateTerminalDeferred: {},
	},
	WorkStateBlocked: {
		WorkStateReady: {}, WorkStateStopped: {},
	},
	WorkStateTerminalDeferred: {
		WorkStateReady: {},
	},
	WorkStateStopped: {
		WorkStateReady: {}, WorkStateReviewing: {}, WorkStatePendingIntegration: {},
	},
	WorkStateCompleted: {},
}

func CanTransition(from, to WorkState) bool {
	if from == to {
		return true
	}
	_, ok := legalWorkTransitions[from][to]
	return ok
}

func transitionWork(record *WorkRecord, to WorkState) error {
	if record == nil {
		return fmt.Errorf("agenttask: nil work record")
	}
	if !CanTransition(record.State, to) {
		return fmt.Errorf("agenttask: illegal work transition %s -> %s", record.State, to)
	}
	record.State = to
	return nil
}

func validateStartRequest(req StartRequest) error {
	fields := []struct {
		name  string
		value string
	}{
		{"command", string(req.CommandID)},
		{"project", string(req.ProjectID)},
		{"workspace", string(req.WorkspaceID)},
		{"milestone", string(req.MilestoneID)},
		{"workflow_revision", string(req.WorkflowRevision)},
		{"config_revision", string(req.ConfigRevision)},
		{"grant_revision", string(req.GrantRevision)},
	}
	for _, field := range fields {
		if err := validateIdentity(field.name, field.value); err != nil {
			return err
		}
	}
	return nil
}

func validateIdentity(field, value string) error {
	if value == "" || strings.TrimSpace(value) != value ||
		strings.ContainsAny(value, "\x00\r\n") {
		return &IdentityError{Field: field, Value: value}
	}
	return nil
}

func validateWorkflowSnapshot(snapshot ProjectWorkflowSnapshot) error {
	if err := validateIdentity("project", string(snapshot.ProjectID)); err != nil {
		return err
	}
	if err := validateIdentity("workspace", string(snapshot.WorkspaceID)); err != nil {
		return err
	}
	if err := validateIdentity("workflow_revision", string(snapshot.Revision)); err != nil {
		return err
	}
	seen := make(map[WorkUnitID]struct{}, len(snapshot.Units))
	for _, unit := range snapshot.Units {
		if err := validateIdentity("work_unit", string(unit.ID)); err != nil {
			return err
		}
		if err := validateIdentity("milestone", string(unit.MilestoneID)); err != nil {
			return err
		}
		if _, ok := seen[unit.ID]; ok {
			return fmt.Errorf("agenttask: duplicate work identity %q", unit.ID)
		}
		seen[unit.ID] = struct{}{}
		switch unit.IsolationMode {
		case agentguardIsolationOverlay, agentguardIsolationWorktree, agentguardIsolationClone:
		default:
			return fmt.Errorf("agenttask: work %q has invalid isolation mode %q", unit.ID, unit.IsolationMode)
		}
	}
	return nil
}

// Local aliases avoid making the transition validator depend on string
// literals while keeping agentguard as the source of isolation mode values.
const (
	agentguardIsolationOverlay  = "overlay"
	agentguardIsolationWorktree = "worktree"
	agentguardIsolationClone    = "clone"
)

func cloneState(state ManagerState) ManagerState {
	out := state
	if out.SchemaVersion == 0 {
		out.SchemaVersion = currentSchemaVersion
	}
	out.Commands = maps.Clone(state.Commands)
	out.Projects = make(map[ProjectID]ProjectRecord, len(state.Projects))
	for id, project := range state.Projects {
		out.Projects[id] = cloneProject(project)
	}
	out.IntegrationLeases = maps.Clone(state.IntegrationLeases)
	out.WorkspaceLeases = maps.Clone(state.WorkspaceLeases)
	if state.DeviceLease != nil {
		lease := *state.DeviceLease
		out.DeviceLease = &lease
	}
	if out.Commands == nil {
		out.Commands = make(map[CommandID]CommandRecord)
	}
	if out.Projects == nil {
		out.Projects = make(map[ProjectID]ProjectRecord)
	}
	if out.IntegrationLeases == nil {
		out.IntegrationLeases = make(map[WorkspaceID]LeaseRecord)
	}
	if out.WorkspaceLeases == nil {
		out.WorkspaceLeases = make(map[WorkspaceID]LeaseRecord)
	}
	return out
}

func cloneProject(project ProjectRecord) ProjectRecord {
	out := project
	if project.Intent != nil {
		intent := *project.Intent
		out.Intent = &intent
	}
	if project.Workflow != nil {
		workflow := cloneWorkflow(*project.Workflow)
		out.Workflow = &workflow
	}
	out.Works = make(map[WorkUnitID]WorkRecord, len(project.Works))
	for id, work := range project.Works {
		out.Works[id] = cloneWork(work)
	}
	if project.Lease != nil {
		lease := *project.Lease
		out.Lease = &lease
	}
	if project.Blocker != nil {
		blocker := *project.Blocker
		out.Blocker = &blocker
	}
	return out
}

func cloneWorkflow(workflow ProjectWorkflowSnapshot) ProjectWorkflowSnapshot {
	out := workflow
	out.Units = make([]WorkUnit, len(workflow.Units))
	for index, unit := range workflow.Units {
		out.Units[index] = cloneUnit(unit)
	}
	return out
}

func cloneUnit(unit WorkUnit) WorkUnit {
	out := unit
	out.Aliases = slices.Clone(unit.Aliases)
	out.ExplicitPredecessors = slices.Clone(unit.ExplicitPredecessors)
	out.DeclaredWriteSet = slices.Clone(unit.DeclaredWriteSet)
	out.Metadata = maps.Clone(unit.Metadata)
	return out
}

func cloneWork(work WorkRecord) WorkRecord {
	out := work
	out.Unit = cloneUnit(work.Unit)
	if work.Target != nil {
		value := *work.Target
		out.Target = &value
	}
	if work.ContinuationTarget != nil {
		value := *work.ContinuationTarget
		out.ContinuationTarget = &value
	}
	if work.Isolation != nil {
		value := *work.Isolation
		out.Isolation = &value
	}
	if work.Submission != nil {
		value := *work.Submission
		value.Metadata = maps.Clone(work.Submission.Metadata)
		value.Locators = slices.Clone(work.Submission.Locators)
		out.Submission = &value
	}
	if work.Review != nil {
		value := *work.Review
		if work.Review.ChangeSet != nil {
			changeSet := *work.Review.ChangeSet
			value.ChangeSet = &changeSet
		}
		out.Review = &value
	}
	if work.ChangeSet != nil {
		value := *work.ChangeSet
		out.ChangeSet = &value
	}
	if work.Integration != nil {
		value := *work.Integration
		if work.Integration.CompletionLocator != nil {
			locator := *work.Integration.CompletionLocator
			value.CompletionLocator = &locator
		}
		if work.Integration.Blocker != nil {
			blocker := *work.Integration.Blocker
			value.Blocker = &blocker
		}
		out.Integration = &value
	}
	if len(work.AttemptObservations) > 0 {
		out.AttemptObservations = make([]AttemptObservationRecord, len(work.AttemptObservations))
		for index, observation := range work.AttemptObservations {
			out.AttemptObservations[index] = observation
			out.AttemptObservations[index].Observation = observation.Observation.Clone()
		}
	}
	if work.Blocker != nil {
		value := *work.Blocker
		out.Blocker = &value
	}
	out.Locators = maps.Clone(work.Locators)
	out.FailureBudgets = maps.Clone(work.FailureBudgets)
	if out.Locators == nil {
		out.Locators = make(map[LocatorKind]LocatorRecord)
	}
	if out.FailureBudgets == nil {
		out.FailureBudgets = make(map[FailureStage]FailureBudgetRecord)
	}
	return out
}

func validateManagerState(state ManagerState) error {
	if state.SchemaVersion != currentSchemaVersion {
		return checkpointError(
			BlockerCorruptCheckpoint,
			fmt.Sprintf("unsupported state schema %d", state.SchemaVersion),
		)
	}
	if err := validateLease("device", state.DeviceLease); err != nil {
		return err
	}
	for workspaceID, lease := range state.WorkspaceLeases {
		if err := validateIdentity("workspace_lease_key", string(workspaceID)); err != nil {
			return checkpointIdentityError(err)
		}
		if err := validateLease("workspace", &lease); err != nil {
			return err
		}
	}
	for workspaceID, lease := range state.IntegrationLeases {
		if err := validateIdentity("integration_lease_key", string(workspaceID)); err != nil {
			return checkpointIdentityError(err)
		}
		if err := validateLease("integration", &lease); err != nil {
			return err
		}
	}
	for commandID, command := range state.Commands {
		if commandID != command.Intent.CommandID {
			return checkpointError(
				BlockerStaleCheckpoint,
				fmt.Sprintf("command map key %q does not match its intent", commandID),
			)
		}
		if err := validateStartIntent(command.Intent); err != nil {
			return err
		}
	}
	for projectID, project := range state.Projects {
		if projectID != project.ProjectID {
			return checkpointError(
				BlockerStaleCheckpoint,
				fmt.Sprintf("project map key %q does not match record %q", projectID, project.ProjectID),
			)
		}
		if !validProjectStatus(project.Status) {
			return checkpointError(
				BlockerCorruptCheckpoint,
				fmt.Sprintf("project %q has unsupported status %q", projectID, project.Status),
			)
		}
		if err := validateIdentity("project", string(project.ProjectID)); err != nil {
			return checkpointIdentityError(err)
		}
		if project.WorkspaceID == "" {
			if project.Intent != nil || project.Workflow != nil || project.Status != ProjectStatusBlocked {
				return checkpointError(
					BlockerStaleCheckpoint,
					fmt.Sprintf("project %q has no workspace identity", projectID),
				)
			}
		} else if err := validateIdentity("workspace", string(project.WorkspaceID)); err != nil {
			return checkpointIdentityError(err)
		}
		if project.Intent != nil {
			if err := validateStartIntent(*project.Intent); err != nil {
				return err
			}
			if project.Intent.ProjectID != project.ProjectID ||
				project.Intent.WorkspaceID != project.WorkspaceID {
				return checkpointError(
					BlockerStaleCheckpoint,
					fmt.Sprintf("project %q intent identity mismatch", projectID),
				)
			}
		}
		if err := validateLease("project", project.Lease); err != nil {
			return err
		}
		for workID, work := range project.Works {
			if workID != work.Unit.ID {
				return checkpointError(
					BlockerStaleCheckpoint,
					fmt.Sprintf("work map key %q does not match record %q", workID, work.Unit.ID),
				)
			}
			if err := validateWorkCheckpoint(project, work); err != nil {
				return err
			}
		}
	}
	return nil
}

func locatorForIsolation(
	project ProjectRecord,
	work WorkRecord,
	isolation IsolationIdentity,
) LocatorRecord {
	return LocatorRecord{
		Kind: LocatorOverlay, Opaque: isolation.ID, Revision: isolation.Revision,
		ProjectID: project.ProjectID, WorkspaceID: project.WorkspaceID,
		WorkUnitID: work.Unit.ID, AttemptID: work.AttemptID,
	}
}

func locatorForChangeSet(
	project ProjectRecord,
	work WorkRecord,
	changeSet ChangeSetIdentity,
) LocatorRecord {
	return LocatorRecord{
		Kind: LocatorChangeSet, Opaque: string(changeSet.ID), Revision: changeSet.Revision,
		ProjectID: project.ProjectID, WorkspaceID: project.WorkspaceID,
		WorkUnitID: work.Unit.ID, AttemptID: work.AttemptID,
	}
}

func failureStageForState(state WorkState) FailureStage {
	switch state {
	case WorkStatePreparing, WorkStateDispatching:
		return FailureStageDispatch
	case WorkStateSubmitted, WorkStateReviewing:
		return FailureStageReview
	case WorkStatePendingIntegration, WorkStateIntegrating:
		return FailureStageIntegration
	default:
		return FailureStageRecovery
	}
}

func (m *Manager) recordFailure(
	work *WorkRecord,
	stage FailureStage,
	blocker Blocker,
) Blocker {
	if work.FailureBudgets == nil {
		work.FailureBudgets = make(map[FailureStage]FailureBudgetRecord)
	}
	budget := work.FailureBudgets[stage]
	if budget.Stage == "" {
		budget = FailureBudgetRecord{
			Stage: stage,
			Limit: m.config.MaxFailureAttempts,
		}
	}
	if budget.Consecutive < budget.Limit {
		budget.Consecutive++
	}
	budget.LastCode = blocker.Code
	budget.AttemptID = work.AttemptID
	budget.UpdatedAt = m.clock.Now()
	work.FailureBudgets[stage] = budget
	if budget.Consecutive >= budget.Limit {
		return Blocker{
			Code: BlockerFailureBudgetExhausted,
			Message: fmt.Sprintf(
				"%s failure budget exhausted after %d consecutive failures; last=%s: %s",
				stage, budget.Consecutive, blocker.Code, blocker.Message,
			),
			Retryable: false,
		}
	}
	return blocker
}

func resetFailure(work *WorkRecord, stage FailureStage) {
	if work.FailureBudgets == nil {
		return
	}
	delete(work.FailureBudgets, stage)
}

func validateStartIntent(intent StartIntent) error {
	for field, value := range map[string]string{
		"command":           string(intent.CommandID),
		"project":           string(intent.ProjectID),
		"workspace":         string(intent.WorkspaceID),
		"milestone":         string(intent.MilestoneID),
		"workflow_revision": string(intent.WorkflowRevision),
		"config_revision":   string(intent.ConfigRevision),
		"grant_revision":    string(intent.GrantRevision),
	} {
		if err := validateIdentity(field, value); err != nil {
			return checkpointIdentityError(err)
		}
	}
	return nil
}

func validateLease(scope string, lease *LeaseRecord) error {
	if lease == nil {
		return nil
	}
	if err := validateIdentity(scope+"_lease_owner", lease.OwnerID); err != nil {
		return checkpointIdentityError(err)
	}
	if err := validateIdentity(scope+"_lease_token", lease.Token); err != nil {
		return checkpointIdentityError(err)
	}
	if lease.ExpiresAt.IsZero() {
		return checkpointError(
			BlockerCorruptCheckpoint,
			fmt.Sprintf("%s lease has no expiry", scope),
		)
	}
	return nil
}

func validateWorkCheckpoint(project ProjectRecord, work WorkRecord) error {
	if err := validateIdentity("work_unit", string(work.Unit.ID)); err != nil {
		return checkpointIdentityError(err)
	}
	if !validWorkState(work.State) {
		return checkpointError(
			BlockerCorruptCheckpoint,
			fmt.Sprintf("work %q has unsupported state %q", work.Unit.ID, work.State),
		)
	}
	if work.Attempt > 0 {
		if work.AttemptID != attemptID(work.Unit.ID, work.Attempt) {
			return checkpointError(
				BlockerStaleCheckpoint,
				fmt.Sprintf("work %q attempt identity does not match its ordinal", work.Unit.ID),
			)
		}
	} else if work.AttemptID != "" {
		return checkpointError(
			BlockerStaleCheckpoint,
			fmt.Sprintf("work %q has an attempt identity without an attempt ordinal", work.Unit.ID),
		)
	}
	for kind, locator := range work.Locators {
		if kind != locator.Kind {
			return checkpointError(
				BlockerAmbiguousCheckpoint,
				fmt.Sprintf("work %q locator key %q does not match kind %q", work.Unit.ID, kind, locator.Kind),
			)
		}
		if err := validateLocator(project, work, locator); err != nil {
			return err
		}
	}
	for stage, budget := range work.FailureBudgets {
		if !validFailureStage(stage) || stage != budget.Stage || budget.Limit == 0 ||
			budget.Consecutive > budget.Limit {
			return checkpointError(
				BlockerCorruptCheckpoint,
				fmt.Sprintf("work %q has invalid %q failure budget", work.Unit.ID, stage),
			)
		}
		if budget.Consecutive > 0 {
			if budget.LastCode == "" {
				return checkpointError(
					BlockerCorruptCheckpoint,
					fmt.Sprintf("work %q has a failure budget without a failure code", work.Unit.ID),
				)
			}
			if err := validateIdentity("failure_attempt", string(budget.AttemptID)); err != nil {
				return checkpointIdentityError(err)
			}
		}
	}
	if work.ContinuationTarget != nil {
		if err := validateTarget(project, *work.ContinuationTarget); err != nil {
			return checkpointIdentityError(err)
		}
	}
	seenObservations := make(map[AttemptID]struct{}, len(work.AttemptObservations))
	for _, observation := range work.AttemptObservations {
		if err := validateIdentity("observation_attempt", string(observation.AttemptID)); err != nil {
			return checkpointIdentityError(err)
		}
		if _, duplicate := seenObservations[observation.AttemptID]; duplicate {
			return checkpointError(
				BlockerCorruptCheckpoint,
				fmt.Sprintf("work %q repeats failure observation for one attempt", work.Unit.ID),
			)
		}
		seenObservations[observation.AttemptID] = struct{}{}
		if err := validateTarget(project, observation.Target); err != nil {
			return checkpointIdentityError(err)
		}
		if !agentpolicy.ValidateAttemptObservation(observation.Observation) {
			return checkpointError(
				BlockerCorruptCheckpoint,
				fmt.Sprintf("work %q has invalid failure observation", work.Unit.ID),
			)
		}
	}
	if work.Submission != nil {
		if work.Submission.ProjectID != project.ProjectID ||
			work.Submission.WorkUnitID != work.Unit.ID ||
			work.Submission.AttemptID != work.AttemptID {
			return checkpointError(
				BlockerStaleCheckpoint,
				fmt.Sprintf("work %q submission identity mismatch", work.Unit.ID),
			)
		}
		for _, locator := range work.Submission.Locators {
			if err := validateLocator(project, work, locator); err != nil {
				return err
			}
		}
	}
	if work.ChangeSet != nil {
		if err := validateIdentity("change_set", string(work.ChangeSet.ID)); err != nil {
			return checkpointIdentityError(err)
		}
		if err := validateIdentity("change_set_revision", work.ChangeSet.Revision); err != nil {
			return checkpointIdentityError(err)
		}
	}
	if locator, ok := work.Locators[LocatorOverlay]; ok {
		if work.Isolation == nil ||
			locator.Opaque != work.Isolation.ID ||
			locator.Revision != work.Isolation.Revision {
			return checkpointError(
				BlockerStaleCheckpoint,
				fmt.Sprintf("work %q overlay locator does not match isolation identity", work.Unit.ID),
			)
		}
	}
	if locator, ok := work.Locators[LocatorChangeSet]; ok {
		if work.ChangeSet == nil ||
			locator.Opaque != string(work.ChangeSet.ID) ||
			locator.Revision != work.ChangeSet.Revision {
			return checkpointError(
				BlockerStaleCheckpoint,
				fmt.Sprintf("work %q change-set locator does not match change-set identity", work.Unit.ID),
			)
		}
	}
	return nil
}

func validProjectStatus(status ProjectStatus) bool {
	switch status {
	case ProjectStatusObserved, ProjectStatusStarted, ProjectStatusRunning,
		ProjectStatusStopped, ProjectStatusBlocked, ProjectStatusCompleted:
		return true
	default:
		return false
	}
}

func validWorkState(state WorkState) bool {
	switch state {
	case WorkStateObserved, WorkStateReady, WorkStatePreparing,
		WorkStateDispatching, WorkStateSubmitted, WorkStateReviewing,
		WorkStatePendingIntegration, WorkStateIntegrating, WorkStateCompleted,
		WorkStateTerminalDeferred, WorkStateBlocked, WorkStateStopped:
		return true
	default:
		return false
	}
}

func validFailureStage(stage FailureStage) bool {
	switch stage {
	case FailureStageDispatch, FailureStageReview,
		FailureStageIntegration, FailureStageRecovery:
		return true
	default:
		return false
	}
}

func validateLocator(
	project ProjectRecord,
	work WorkRecord,
	locator LocatorRecord,
) error {
	switch locator.Kind {
	case LocatorProcess, LocatorSession, LocatorOverlay, LocatorChangeSet, LocatorCompletion:
	default:
		return checkpointError(
			BlockerCorruptCheckpoint,
			fmt.Sprintf("work %q has unsupported locator kind %q", work.Unit.ID, locator.Kind),
		)
	}
	if locator.ProjectID != project.ProjectID ||
		locator.WorkspaceID != project.WorkspaceID ||
		locator.WorkUnitID != work.Unit.ID ||
		locator.AttemptID != work.AttemptID {
		return checkpointError(
			BlockerStaleCheckpoint,
			fmt.Sprintf("work %q %s locator identity mismatch", work.Unit.ID, locator.Kind),
		)
	}
	if err := validateIdentity(string(locator.Kind)+"_locator", locator.Opaque); err != nil {
		return checkpointIdentityError(err)
	}
	if err := validateIdentity(string(locator.Kind)+"_locator_revision", locator.Revision); err != nil {
		return checkpointIdentityError(err)
	}
	return nil
}

func checkpointIdentityError(err error) error {
	return checkpointError(BlockerStaleCheckpoint, err.Error())
}

func checkpointError(code BlockerCode, message string) error {
	return &CheckpointError{Code: code, Message: message}
}