iop/apps/agent/internal/taskloop/recovery.go

package taskloop

import (
	"context"
	"encoding/json"
	"errors"
	"os"
	"strings"
	"syscall"

	"iop/packages/go/agentconfig"
	"iop/packages/go/agenttask"
)

// Recovery resolves only exact host-owned locators. Malformed or unverifiable
// identity is reported as ambiguous rather than rebound to a current process.
type Recovery struct {
	snapshot agentconfig.RuntimeSnapshot
	roots    ArtifactRootResolver
}

var _ agenttask.RecoveryInspector = (*Recovery)(nil)

func NewRecovery(
	snapshot agentconfig.RuntimeSnapshot,
	roots ArtifactRootResolver,
) *Recovery {
	return &Recovery{snapshot: snapshot, roots: roots}
}

func (recovery *Recovery) Inspect(
	ctx context.Context,
	request agenttask.RecoveryRequest,
) (agenttask.RecoveryObservation, error) {
	if err := ctx.Err(); err != nil {
		return agenttask.RecoveryObservation{}, err
	}
	observation := agenttask.RecoveryObservation{
		ProjectID:   request.Project.ProjectID,
		WorkspaceID: request.Project.WorkspaceID,
		WorkUnitID:  request.Work.Unit.ID,
		AttemptID:   request.Work.AttemptID,
		Execution:   agenttask.RecoveryExecutionAbsent,
		Completion:  agenttask.RecoveryCompletionUnknown,
	}
	if completion, ok := request.Locators[agenttask.LocatorCompletion]; ok {
		switch {
		case request.Work.Integration != nil &&
			request.Work.Integration.Outcome == agenttask.IntegrationOutcomeIntegrated &&
			request.Work.Integration.CompletionLocator != nil &&
			request.Work.Integration.CompletionLocator.Opaque == completion.Opaque &&
			request.Work.Integration.CompletionLocator.Revision == completion.Revision:
			observation.Completion = agenttask.RecoveryCompletionComplete
		case strings.HasPrefix(completion.Opaque, "integration:"):
			observation.Completion = agenttask.RecoveryCompletionPartial
		default:
			observation.Completion = agenttask.RecoveryCompletionAmbiguous
		}
	}
	if request.Work.Submission != nil {
		submission := *request.Work.Submission
		submission.Metadata = cloneStringMap(submission.Metadata)
		submission.Locators = append([]agenttask.LocatorRecord(nil), submission.Locators...)
		observation.Execution = agenttask.RecoveryExecutionSubmitted
		observation.Submission = &submission
		return observation, nil
	}
	process, ok := request.Locators[agenttask.LocatorProcess]
	if !ok {
		return observation, nil
	}
	var reference processReference
	if err := json.Unmarshal([]byte(process.Opaque), &reference); err != nil ||
		reference.PID <= 0 || reference.StartToken == "" {
		observation.Execution = agenttask.RecoveryExecutionAmbiguous
		return observation, nil
	}
	currentToken, tokenErr := processStartToken(reference.PID)
	if tokenErr == nil {
		if currentToken == reference.StartToken {
			observation.Execution = agenttask.RecoveryExecutionLive
		} else {
			observation.Execution = agenttask.RecoveryExecutionAmbiguous
		}
		return observation, nil
	}
	if processMissing(reference.PID) {
		submission, complete, err := recovery.recoverSubmission(request)
		if err != nil {
			return agenttask.RecoveryObservation{}, err
		}
		if complete {
			observation.Execution = agenttask.RecoveryExecutionSubmitted
			observation.Submission = &submission
		} else {
			observation.Execution = agenttask.RecoveryExecutionExited
		}
		return observation, nil
	}
	observation.Execution = agenttask.RecoveryExecutionAmbiguous
	return observation, nil
}

func (recovery *Recovery) recoverSubmission(
	request agenttask.RecoveryRequest,
) (agenttask.Submission, bool, error) {
	registration, ok := recovery.snapshot.Project(string(request.Project.ProjectID))
	if !ok {
		return agenttask.Submission{}, false, errors.New("taskloop: recovery project is no longer registered")
	}
	canonicalRoot, err := canonicalDirectory(registration.Workspace)
	if err != nil {
		return agenttask.Submission{}, false, err
	}
	if WorkspaceIdentity(canonicalRoot) != request.Project.WorkspaceID {
		return agenttask.Submission{}, false, errors.New("taskloop: recovery workspace identity mismatch")
	}
	if recovery.roots == nil {
		return agenttask.Submission{}, false, errors.New("taskloop: retained artifact root resolver is unavailable")
	}
	artifactRoot, err := recovery.roots.ArtifactRoot(request.Work)
	if err != nil {
		return agenttask.Submission{}, false, err
	}
	reviewPath := request.Work.Unit.Metadata["review_path"]
	if reviewPath == "" {
		return agenttask.Submission{}, false, nil
	}
	artifact, err := inspectReviewArtifact(artifactRoot, reviewPath)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return agenttask.Submission{}, false, nil
		}
		return agenttask.Submission{}, false, err
	}
	if artifact.placeholder {
		return agenttask.Submission{}, false, nil
	}
	locators := make([]agenttask.LocatorRecord, 0, len(request.Locators))
	for _, locator := range request.Locators {
		if locator.Kind != agenttask.LocatorProcess &&
			locator.Kind != agenttask.LocatorSession {
			continue
		}
		locators = append(locators, locator)
	}
	return agenttask.Submission{
		ProjectID:  request.Project.ProjectID,
		WorkUnitID: request.Work.Unit.ID,
		AttemptID:  request.Work.AttemptID,
		ArtifactID: artifactIdentity(request.Work),
		Ready:      true,
		Locators:   locators,
	}, true, nil
}

func processMissing(pid int) bool {
	process, err := os.FindProcess(pid)
	if err != nil {
		return true
	}
	err = process.Signal(syscall.Signal(0))
	if err == nil {
		return false
	}
	return errors.Is(err, os.ErrProcessDone) ||
		strings.Contains(strings.ToLower(err.Error()), "no such process")
}

func cloneStringMap(input map[string]string) map[string]string {
	if input == nil {
		return nil
	}
	output := make(map[string]string, len(input))
	for key, value := range input {
		output[key] = value
	}
	return output
}