iop/packages/go/agenttask/state_machine.go

package agenttask

import (
	"fmt"
	"maps"
	"slices"
	"strings"
)

const currentSchemaVersion uint32 = 1

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)
	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)
	}
	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.Isolation != nil {
		value := *work.Isolation
		out.Isolation = &value
	}
	if work.Submission != nil {
		value := *work.Submission
		value.Metadata = maps.Clone(work.Submission.Metadata)
		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.Blocker != nil {
			blocker := *work.Integration.Blocker
			value.Blocker = &blocker
		}
		out.Integration = &value
	}
	if work.Blocker != nil {
		value := *work.Blocker
		out.Blocker = &value
	}
	return out
}