iop/packages/go/agenttask/dispatch.go

package agenttask

import (
	"context"
	"fmt"

	"iop/packages/go/agentguard"
)

func (m *Manager) runWork(
	ctx context.Context,
	projectID ProjectID,
	workID WorkUnitID,
) error {
	for {
		project, work, err := m.loadWork(ctx, projectID, workID)
		if err != nil {
			return err
		}
		switch work.State {
		case WorkStateReviewing:
			if work.Submission == nil {
				m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
					Code:    BlockerArtifactMismatch,
					Message: "reviewing work has no durable submission identity",
				})
				return nil
			}
			rework, reviewErr := m.reviewSubmission(ctx, project, work, *work.Submission)
			if reviewErr != nil || !rework {
				return reviewErr
			}
			continue
		case WorkStateReady:
		default:
			return nil
		}

		if err := m.changeWork(ctx, projectID, workID, func(work *WorkRecord) error {
			return transitionWork(work, WorkStatePreparing)
		}); err != nil {
			return err
		}
		project, work, err = m.loadWork(ctx, projectID, workID)
		if err != nil {
			return err
		}
		target, err := m.selector.Select(ctx, SelectionRequest{Project: project, Work: work})
		if err != nil {
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code: BlockerSelectionFailed, Message: err.Error(), Retryable: true,
			})
			return nil
		}
		if err := validateTarget(project, target); err != nil {
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code: BlockerInvalidIdentity, Message: err.Error(),
			})
			return nil
		}
		isolation, err := m.isolation.Prepare(ctx, IsolationRequest{
			Project: project, Work: work, Target: target,
			IdempotencyKey: dispatchKey(projectID, workID, work.Attempt) + "/isolation",
		})
		if err != nil {
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code: BlockerIsolationFailed, Message: err.Error(), Retryable: true,
			})
			return nil
		}
		admissionRequest, isolationIdentity, err := validatePreparedIsolation(project, work, target, isolation)
		if err != nil {
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code: BlockerIsolationFailed, Message: err.Error(),
			})
			return nil
		}
		admission := agentguard.Admit(admissionRequest)
		if !admission.Allowed() {
			detail := "workspace admission denied"
			if admission.Blocker != nil {
				detail = string(admission.Blocker.Code) + ": " + admission.Blocker.Message
			}
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code: BlockerAdmissionFailed, Message: detail,
			})
			return nil
		}
		lease, err := m.scheduler.Acquire(ctx, DispatchCandidate{
			ProjectID: projectID, WorkspaceID: project.WorkspaceID,
			WorkUnitID: workID, AttemptID: work.AttemptID, Target: target,
		})
		if err != nil {
			if ctx.Err() != nil {
				_ = m.changeWork(context.WithoutCancel(ctx), projectID, workID, func(work *WorkRecord) error {
					return transitionWork(work, WorkStateReady)
				})
				return ctx.Err()
			}
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code: BlockerProviderCapacity, Message: err.Error(), Retryable: true,
			})
			return nil
		}
		err = m.changeWork(ctx, projectID, workID, func(work *WorkRecord) error {
			if err := transitionWork(work, WorkStateDispatching); err != nil {
				return err
			}
			work.Target = &target
			work.Isolation = &isolationIdentity
			return nil
		})
		if err != nil {
			lease.Release()
			return err
		}
		project, work, err = m.loadWork(ctx, projectID, workID)
		if err != nil {
			lease.Release()
			return err
		}
		var cmdID CommandID
		var wfRev WorkflowRevision
		if project.Intent != nil {
			cmdID = project.Intent.CommandID
			wfRev = project.Intent.WorkflowRevision
		}
		m.emit(ctx, Event{
			Type: EventDispatchStarted, ProjectID: projectID, WorkspaceID: project.WorkspaceID,
			WorkUnitID: workID, CommandID: cmdID, WorkflowRevision: wfRev, AttemptID: work.AttemptID, Ordinal: work.DispatchOrdinal,
			State: work.State, ProviderID: target.ProviderID, ProfileID: target.ProfileID,
			WriteSetKind: work.Unit.WriteSetKind, IsolationMode: work.Unit.IsolationMode,
		})
		var submission Submission
		dispatchRequest := DispatchRequest{
			Project: project, Work: work, Target: target, AdmissionRequest: admissionRequest,
			Permit: admission.Permit, Workspace: *admission.Workspace,
			IdempotencyKey: dispatchKey(projectID, workID, work.Attempt),
		}
		validation, invokeErr := agentguard.Invoke(
			ctx,
			admission.Permit,
			admissionRequest,
			func(invokeCtx context.Context, workspace agentguard.CanonicalWorkspace) error {
				dispatchRequest.Workspace = workspace
				var err error
				submission, err = m.invoker.Invoke(invokeCtx, dispatchRequest)
				return err
			},
		)
		lease.Release()
		if !validation.Allowed() {
			detail := "admission permit became invalid before invocation"
			if validation.Blocker != nil {
				detail = string(validation.Blocker.Code) + ": " + validation.Blocker.Message
			}
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code: BlockerAdmissionFailed, Message: detail,
			})
			return nil
		}
		if invokeErr != nil {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code: BlockerInvocationFailed, Message: invokeErr.Error(), Retryable: true,
			})
			return nil
		}
		if err := validateSubmission(projectID, work, submission); err != nil {
			m.blockWork(ctx, projectID, workID, WorkStateBlocked, Blocker{
				Code: BlockerArtifactMismatch, 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 err := m.changeWork(ctx, projectID, workID, func(work *WorkRecord) error {
			if err := transitionWork(work, WorkStateSubmitted); err != nil {
				return err
			}
			work.Submission = &submission
			return nil
		}); err != nil {
			return err
		}
		m.emit(ctx, Event{
			Type: EventSubmissionAccepted, ProjectID: projectID, WorkspaceID: project.WorkspaceID,
			WorkUnitID: workID, CommandID: cmdID, WorkflowRevision: wfRev, AttemptID: work.AttemptID, Ordinal: work.DispatchOrdinal,
			Detail: string(submission.ArtifactID),
		})
		if err := m.changeWork(ctx, projectID, workID, func(work *WorkRecord) error {
			return transitionWork(work, WorkStateReviewing)
		}); err != nil {
			return err
		}
		project, work, err = m.loadWork(ctx, projectID, workID)
		if err != nil {
			return err
		}
		rework, err := m.reviewSubmission(ctx, project, work, submission)
		if err != nil || !rework {
			return err
		}
	}
}

func validateTarget(project ProjectRecord, target ExecutionTarget) error {
	for field, value := range map[string]string{
		"provider":         target.ProviderID,
		"model":            target.ModelID,
		"profile":          target.ProfileID,
		"profile_revision": target.ProfileRevision,
	} {
		if err := validateIdentity(field, value); err != nil {
			return err
		}
	}
	if project.Intent == nil || target.ConfigRevision != project.Intent.ConfigRevision {
		return fmt.Errorf("agenttask: selected target config revision does not match manual start")
	}
	if target.Capacity <= 0 {
		return fmt.Errorf("agenttask: selected target capacity must be positive")
	}
	return nil
}

func validatePreparedIsolation(
	project ProjectRecord,
	work WorkRecord,
	target ExecutionTarget,
	prepared PreparedIsolation,
) (agentguard.AdmissionRequest, IsolationIdentity, error) {
	if prepared.Grant == nil || prepared.Descriptor == nil {
		return agentguard.AdmissionRequest{}, IsolationIdentity{},
			fmt.Errorf("agenttask: isolation backend returned incomplete strict ports")
	}
	if project.Intent == nil {
		return agentguard.AdmissionRequest{}, IsolationIdentity{},
			fmt.Errorf("agenttask: project has no manual start intent")
	}
	if prepared.Grant.ProjectID != string(project.ProjectID) ||
		prepared.Grant.WorkspaceID != string(project.WorkspaceID) ||
		prepared.Grant.Revision != string(project.Intent.GrantRevision) {
		return agentguard.AdmissionRequest{}, IsolationIdentity{},
			fmt.Errorf("agenttask: workspace grant identity or revision mismatch")
	}
	if prepared.Descriptor.Mode != work.Unit.IsolationMode {
		return agentguard.AdmissionRequest{}, IsolationIdentity{},
			fmt.Errorf("agenttask: isolation mode differs from the workflow request")
	}
	if prepared.Profile.ProviderID != target.ProviderID ||
		prepared.Profile.ModelID != target.ModelID ||
		prepared.Profile.ProfileID != target.ProfileID ||
		prepared.Profile.Revision != target.ProfileRevision {
		return agentguard.AdmissionRequest{}, IsolationIdentity{},
			fmt.Errorf("agenttask: admitted provider profile differs from selected target")
	}
	identity := IsolationIdentity{
		ID: prepared.Descriptor.ID, Revision: prepared.Descriptor.Revision,
		Mode: prepared.Descriptor.Mode, PinnedBaseRevision: prepared.Descriptor.PinnedBaseRevision,
		TaskRoot: prepared.Descriptor.TaskRoot,
	}
	for field, value := range map[string]string{
		"isolation":            identity.ID,
		"isolation_revision":   identity.Revision,
		"pinned_base_revision": identity.PinnedBaseRevision,
	} {
		if err := validateIdentity(field, value); err != nil {
			return agentguard.AdmissionRequest{}, IsolationIdentity{}, err
		}
	}
	return agentguard.AdmissionRequest{
		Grant: prepared.Grant, Isolation: prepared.Descriptor, Profile: prepared.Profile,
	}, identity, nil
}

func validateSubmission(projectID ProjectID, work WorkRecord, submission Submission) error {
	if submission.ProjectID != projectID || submission.WorkUnitID != work.Unit.ID ||
		submission.AttemptID != work.AttemptID {
		return fmt.Errorf("agenttask: submission durable identity mismatch")
	}
	if err := validateIdentity("artifact", string(submission.ArtifactID)); err != nil {
		return err
	}
	return nil
}

func (m *Manager) loadWork(
	ctx context.Context,
	projectID ProjectID,
	workID WorkUnitID,
) (ProjectRecord, WorkRecord, error) {
	state, err := m.load(ctx)
	if err != nil {
		return ProjectRecord{}, WorkRecord{}, err
	}
	project, ok := state.Projects[projectID]
	if !ok {
		return ProjectRecord{}, WorkRecord{}, fmt.Errorf("agenttask: project %q disappeared", projectID)
	}
	work, ok := project.Works[workID]
	if !ok {
		return ProjectRecord{}, WorkRecord{}, fmt.Errorf("agenttask: work %q disappeared", workID)
	}
	return project, work, nil
}

func (m *Manager) changeWork(
	ctx context.Context,
	projectID ProjectID,
	workID WorkUnitID,
	change func(*WorkRecord) error,
) error {
	return m.mutate(ctx, func(state *ManagerState) error {
		project, ok := state.Projects[projectID]
		if !ok {
			return fmt.Errorf("agenttask: project %q disappeared", projectID)
		}
		work, ok := project.Works[workID]
		if !ok {
			return fmt.Errorf("agenttask: work %q disappeared", workID)
		}
		if err := change(&work); err != nil {
			return err
		}
		work.UpdatedAt = m.clock.Now()
		project.Works[workID] = work
		project.UpdatedAt = m.clock.Now()
		state.Projects[projectID] = project
		return nil
	})
}

func (m *Manager) blockWork(
	ctx context.Context,
	projectID ProjectID,
	workID WorkUnitID,
	state WorkState,
	blocker Blocker,
) {
	_ = m.changeWork(ctx, projectID, workID, func(work *WorkRecord) error {
		if CanTransition(work.State, state) {
			work.State = state
		} else {
			work.State = WorkStateBlocked
		}
		work.Blocker = &blocker
		return nil
	})
	m.emit(ctx, Event{
		Type: EventBlocked, ProjectID: projectID, WorkUnitID: workID,
		State: state, Detail: string(blocker.Code),
	})
}