gito/services/core/internal/worker/runner.go

package worker

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"
	"time"

	"git.toki-labs.com/toki/gito/services/core/internal/config"
	"git.toki-labs.com/toki/gito/services/core/internal/core"
	"git.toki-labs.com/toki/gito/services/core/internal/events"
	"git.toki-labs.com/toki/gito/services/core/internal/gitengine"
	"git.toki-labs.com/toki/gito/services/core/internal/storage"
)

type OperationPicker interface {
	PickQueuedOperation(ctx context.Context, now time.Time) (core.Operation, bool, error)
}

type Clock interface {
	Now() time.Time
}

type realClock struct{}

func (realClock) Now() time.Time {
	return time.Now()
}

type AgentRunInput struct {
	OperationID      string
	RepoID           string
	Branch           string
	WorkspacePath    string
	Instruction      string
	PolicyContext    map[string]any
	ExpectedRevision string
	CredentialRefs   map[string]string
}

type AgentRunState string

const (
	AgentRunStateSucceeded AgentRunState = "succeeded"
	AgentRunStateFailed    AgentRunState = "failed"
	AgentRunStateCancelled AgentRunState = "cancelled"
)

type AgentRunResult struct {
	State   AgentRunState
	Message string
}

type AgentRunGitEvidence struct {
	StatusClean    bool
	StatusOutput   string
	BeforeRevision string
	AfterRevision  string
	ChangedFiles   []gitengine.ChangedFile
}

type RevisionScanRequest struct {
	Provider   string
	RepoID     string
	Branch     string
	WorkDir    string
	Runner     gitengine.CommandRunner
	ObservedAt time.Time
}

type RevisionScanResult struct {
	Attempted bool
	Matched   bool
	EventID   string
}

type RevisionScanner interface {
	ScanAgentRunRevision(ctx context.Context, req RevisionScanRequest) (RevisionScanResult, error)
}

type AgentRunInputLoader interface {
	LoadAgentRunInput(ctx context.Context, op core.Operation) (AgentRunInput, error)
}

type AgentRunInvoker interface {
	InvokeAgentRun(ctx context.Context, input AgentRunInput) (AgentRunResult, error)
}

type RunnerDependencies struct {
	OperationStore      storage.OperationStore
	OperationEventStore storage.OperationEventStore
	AgentRunInputLoader AgentRunInputLoader
	AgentRunInvoker     AgentRunInvoker
	GitRunner           gitengine.CommandRunner
	RevisionScanner     RevisionScanner
}

type Runner struct {
	cfg                 config.Config
	logger              *slog.Logger
	picker              OperationPicker
	clock               Clock
	operationStore      storage.OperationStore
	operationEventStore storage.OperationEventStore
	agentRunInputLoader AgentRunInputLoader
	agentRunInvoker     AgentRunInvoker
	gitRunner           gitengine.CommandRunner
	revisionScanner     RevisionScanner
}

func NewRunner(cfg config.Config, logger *slog.Logger, picker OperationPicker, clock Clock, scanner RevisionScanner) *Runner {
	return NewRunnerWithDependencies(cfg, logger, picker, clock, RunnerDependencies{
		RevisionScanner: scanner,
	})
}

func NewRunnerWithDependencies(cfg config.Config, logger *slog.Logger, picker OperationPicker, clock Clock, deps RunnerDependencies) *Runner {
	if clock == nil {
		clock = realClock{}
	}
	if logger == nil {
		logger = slog.Default()
	}
	gitRunner := deps.GitRunner
	if gitRunner == nil {
		gitRunner = gitengine.CLI{}
	}
	return &Runner{
		cfg:                 cfg,
		logger:              logger,
		picker:              picker,
		clock:               clock,
		operationStore:      deps.OperationStore,
		operationEventStore: deps.OperationEventStore,
		agentRunInputLoader: deps.AgentRunInputLoader,
		agentRunInvoker:     deps.AgentRunInvoker,
		gitRunner:           gitRunner,
		revisionScanner:     deps.RevisionScanner,
	}
}

func (r *Runner) Run() error {
	return r.RunOnce(context.Background())
}

func (r *Runner) RunOnce(ctx context.Context) error {
	if !r.cfg.WorkerEnabled {
		r.logger.Info("worker disabled")
		return nil
	}

	if r.picker == nil {
		r.logger.Warn("worker picker not configured")
		return nil
	}

	op, ok, err := r.picker.PickQueuedOperation(ctx, r.clock.Now())
	if err != nil {
		r.logger.Error("failed to pick operation", "error", err)
		return fmt.Errorf("pick operation: %w", err)
	}

	if !ok {
		r.logger.Debug("no pending operation found")
		return nil
	}

	r.logger.Info("picked operation", "operation_id", op.ID, "type", op.Type)
	if op.Type == core.OperationAgentRun {
		return r.runAgentRunOperation(ctx, op)
	}
	return nil
}

func (r *Runner) runAgentRunOperation(ctx context.Context, op core.Operation) error {
	if r.operationStore == nil || r.operationEventStore == nil || r.agentRunInputLoader == nil || r.agentRunInvoker == nil {
		return fmt.Errorf("agent_run handler is not fully configured")
	}

	startedAt := r.clock.Now().UTC()
	started := storage.OperationEvent{
		OperationID: op.ID,
		Event: events.Event{
			ID:      eventID(op.ID, events.OperationStarted, startedAt),
			Type:    events.OperationStarted,
			Subject: operationEventSubject(op.ID),
		},
	}
	startPayload, err := agentRunEventPayload(op, core.OperationRunning, startedAt, nil)
	if err != nil {
		return err
	}
	started.Event.Payload = startPayload
	started.Event.CreatedAt = startedAt

	if err := r.operationEventStore.AppendEvent(ctx, started); err != nil {
		return fmt.Errorf("append agent_run started event: %w", err)
	}

	agentStarted := storage.OperationEvent{
		OperationID: op.ID,
		Event: events.Event{
			ID:      eventID(op.ID, events.AgentRunStarted, startedAt),
			Type:    events.AgentRunStarted,
			Subject: operationEventSubject(op.ID),
		},
	}
	agentRunStartedPayload, err := agentRunEventPayload(op, core.OperationRunning, startedAt, nil)
	if err != nil {
		return err
	}
	agentStarted.Event.Payload = agentRunStartedPayload
	agentStarted.Event.CreatedAt = startedAt

	if err := r.operationEventStore.AppendEvent(ctx, agentStarted); err != nil {
		return fmt.Errorf("append agent_run_started event: %w", err)
	}

	input, err := r.agentRunInputLoader.LoadAgentRunInput(ctx, op)
	if err != nil {
		return fmt.Errorf("load agent_run input: %w", err)
	}

	result, err := r.agentRunInvoker.InvokeAgentRun(ctx, input)
	terminalState := core.OperationFailed
	completionMessage := strings.TrimSpace(result.Message)
	// If backend invocation fails, fail the operation with a stable, sanitized message.
	if err != nil {
		completionMessage = "agent run invocation failed"
	} else {
		terminalState, err = toTerminalOperationState(result.State)
		if err != nil {
			return fmt.Errorf("unsupported agent_run result state: %w", err)
		}
	}
	if completionMessage == "" {
		completionMessage = string(terminalState)
	}

	policy := agentRunGitResultPolicyFrom(input.PolicyContext, input.Branch)
	completionExtra := map[string]any{"message": completionMessage}
	if terminalState == core.OperationSucceeded {
		evidence, evidenceErr := r.collectAgentRunGitEvidence(input)
		if evidenceErr != nil {
			r.logger.Warn("agent_run git evidence collection failed", "operation_id", op.ID, "error", evidenceErr)
			terminalState = core.OperationFailed
			completionMessage = "git evidence collection failed"
			completionExtra = map[string]any{"message": completionMessage}
		} else {
			completionExtra, terminalState, completionMessage = r.applyAgentRunGitResultPolicy(input, policy, evidence, completionMessage)
			if terminalState == core.OperationSucceeded && policy.CommitPushAllowed && r.revisionScanner != nil {
				var provider string
				if pVal, ok := input.PolicyContext["provider"]; ok {
					if s, ok := pVal.(string); ok {
						provider = strings.TrimSpace(s)
					}
				}
				scanReq := RevisionScanRequest{
					Provider:   provider,
					RepoID:     op.RepoID,
					Branch:     policy.Branch,
					WorkDir:    input.WorkspacePath,
					Runner:     r.gitRunner,
					ObservedAt: r.clock.Now().UTC(),
				}
				scanRes, err := r.revisionScanner.ScanAgentRunRevision(ctx, scanReq)
				if err != nil {
					r.logger.Warn("agent_run revision scan failed", "operation_id", op.ID, "error", err)
					terminalState = core.OperationFailed
					completionMessage = "revision scan failed"
					completionExtra["message"] = completionMessage
					completionExtra["revision_scan_attempted"] = true
					completionExtra["revision_event_matched"] = false
				} else {
					completionExtra["revision_scan_attempted"] = scanRes.Attempted
					completionExtra["revision_event_matched"] = scanRes.Matched
					if scanRes.EventID != "" {
						completionExtra["revision_event_id"] = scanRes.EventID
					}
				}
			}
		}
	}

	terminalAt := r.clock.Now().UTC()
	var updatedOperation core.Operation
	switch terminalState {
	case core.OperationSucceeded:
		updatedOperation, err = r.operationStore.SucceedOperation(ctx, op.ID, terminalAt)
	case core.OperationFailed:
		updatedOperation, err = r.operationStore.FailOperation(ctx, op.ID, terminalAt)
	case core.OperationCancelled:
		updatedOperation, err = r.operationStore.CancelOperation(ctx, op.ID, terminalAt)
	default:
		err = fmt.Errorf("unsupported terminal state %q", terminalState)
	}
	if err != nil {
		return fmt.Errorf("mark agent_run operation terminal: %w", err)
	}

	completionEventType := events.OperationCompleted
	if terminalState == core.OperationFailed {
		completionEventType = events.OperationFailed
	}
	completedAt := storage.OperationEvent{
		OperationID: op.ID,
		Event: events.Event{
			ID:      eventID(op.ID, completionEventType, terminalAt),
			Type:    completionEventType,
			Subject: operationEventSubject(op.ID),
		},
	}
	completedPayload, err := agentRunEventPayload(updatedOperation, terminalState, terminalAt, completionExtra)
	if err != nil {
		return err
	}
	completedAt.Event.Payload = completedPayload
	completedAt.Event.CreatedAt = terminalAt
	if err := r.operationEventStore.AppendEvent(ctx, completedAt); err != nil {
		return fmt.Errorf("append agent_run completion event: %w", err)
	}

	agentCompleted := storage.OperationEvent{
		OperationID: op.ID,
		Event: events.Event{
			ID:      eventID(op.ID, events.AgentRunCompleted, terminalAt),
			Type:    events.AgentRunCompleted,
			Subject: operationEventSubject(op.ID),
		},
	}
	agentCompletedPayload, err := agentRunEventPayload(updatedOperation, terminalState, terminalAt, completionExtra)
	if err != nil {
		return err
	}
	agentCompleted.Event.Payload = agentCompletedPayload
	agentCompleted.Event.CreatedAt = terminalAt
	if err := r.operationEventStore.AppendEvent(ctx, agentCompleted); err != nil {
		return fmt.Errorf("append agent_run_completed event: %w", err)
	}

	return nil
}

func (r *Runner) collectAgentRunGitEvidence(input AgentRunInput) (AgentRunGitEvidence, error) {
	workdir := strings.TrimSpace(input.WorkspacePath)
	before := strings.TrimSpace(input.ExpectedRevision)
	if workdir == "" || before == "" {
		return AgentRunGitEvidence{}, fmt.Errorf("workspace_path and expected_revision are required")
	}

	status, err := gitengine.Status(r.gitRunner, workdir)
	if err != nil {
		return AgentRunGitEvidence{}, fmt.Errorf("git status: %w", err)
	}
	after, err := gitengine.HeadRevision(r.gitRunner, workdir)
	if err != nil {
		return AgentRunGitEvidence{}, fmt.Errorf("git head revision: %w", err)
	}
	changedFiles, err := gitengine.ChangedFilesWithStatus(r.gitRunner, workdir, before, after)
	if err != nil {
		return AgentRunGitEvidence{}, fmt.Errorf("git changed files: %w", err)
	}

	return AgentRunGitEvidence{
		StatusClean:    status.Clean,
		StatusOutput:   status.Output,
		BeforeRevision: before,
		AfterRevision:  after,
		ChangedFiles:   changedFiles,
	}, nil
}

func agentRunGitEvidencePayload(evidence AgentRunGitEvidence) map[string]any {
	changedFiles := make([]any, 0, len(evidence.ChangedFiles))
	for _, file := range evidence.ChangedFiles {
		changedFiles = append(changedFiles, map[string]any{
			"path":        file.Path,
			"change_type": file.ChangeType,
		})
	}
	return map[string]any{
		"status_clean":    evidence.StatusClean,
		"status_output":   evidence.StatusOutput,
		"before_revision": evidence.BeforeRevision,
		"after_revision":  evidence.AfterRevision,
		"changed_files":   changedFiles,
	}
}

func operationEventSubject(operationID string) string {
	return "operation:" + operationID
}

func eventID(operationID, eventType string, at time.Time) string {
	return fmt.Sprintf("agent-run:%s:%s:%d", operationID, eventType, at.UnixNano())
}

func agentRunEventPayload(op core.Operation, state core.OperationState, at time.Time, extra map[string]any) ([]byte, error) {
	payload := map[string]any{
		"operation_id": op.ID,
		"repo_id":      op.RepoID,
		"type":         op.Type,
		"state":        state,
		"at":           at.Format(time.RFC3339Nano),
	}
	for key, value := range extra {
		payload[key] = value
	}
	encoded, err := json.Marshal(payload)
	if err != nil {
		return nil, fmt.Errorf("serialize agent run event payload: %w", err)
	}
	return encoded, nil
}

type agentRunGitResultPolicy struct {
	CommitPushAllowed bool
	CommitMessage     string
	Remote            string
	Branch            string
}

func agentRunGitResultPolicyFrom(ctx map[string]any, branch string) agentRunGitResultPolicy {
	p := agentRunGitResultPolicy{
		Remote: "origin",
		Branch: strings.TrimSpace(branch),
	}
	if ctx == nil {
		return p
	}
	if v, ok := ctx["commit_push"]; ok {
		if b, ok := v.(bool); ok {
			p.CommitPushAllowed = b
		}
	}
	if !p.CommitPushAllowed {
		if v, ok := ctx["allowed_result_actions"]; ok {
			if agentRunAllowsCommitPushAction(v) {
				p.CommitPushAllowed = true
			}
		}
	}
	if msg, ok := ctx["commit_message"]; ok {
		if s, ok := msg.(string); ok {
			p.CommitMessage = strings.TrimSpace(s)
		}
	}
	if remote, ok := ctx["remote"]; ok {
		if s, ok := remote.(string); ok && strings.TrimSpace(s) != "" {
			p.Remote = strings.TrimSpace(s)
		}
	}
	if b, ok := ctx["branch"]; ok {
		if s, ok := b.(string); ok && strings.TrimSpace(s) != "" {
			p.Branch = strings.TrimSpace(s)
		}
	}
	return p
}

func agentRunAllowsCommitPushAction(value any) bool {
	switch actions := value.(type) {
	case []any:
		for _, action := range actions {
			if s, ok := action.(string); ok && strings.TrimSpace(s) == "commit_push" {
				return true
			}
		}
	case []string:
		for _, action := range actions {
			if strings.TrimSpace(action) == "commit_push" {
				return true
			}
		}
	}
	return false
}

func (r *Runner) applyAgentRunGitResultPolicy(
	input AgentRunInput,
	policy agentRunGitResultPolicy,
	evidence AgentRunGitEvidence,
	completionMessage string,
) (map[string]any, core.OperationState, string) {
	extra := agentRunGitEvidencePayload(evidence)
	extra["commit_push_allowed"] = policy.CommitPushAllowed

	if !policy.CommitPushAllowed {
		if !evidence.StatusClean {
			msg := "workspace dirty after agent run"
			extra["message"] = msg
			return extra, core.OperationFailed, msg
		}
		extra["message"] = completionMessage
		return extra, core.OperationSucceeded, completionMessage
	}

	commitMsg := policy.CommitMessage
	if commitMsg == "" {
		commitMsg = "agent run result: " + input.OperationID
	}
	extra["commit_attempted"] = true
	if err := gitengine.Commit(r.gitRunner, input.WorkspacePath, gitengine.CommitOptions{
		Message: commitMsg,
		All:     true,
	}); err != nil {
		msg := "git commit failed"
		extra["message"] = msg
		return extra, core.OperationFailed, msg
	}
	extra["push_attempted"] = true
	extra["push_remote"] = policy.Remote
	extra["push_branch"] = policy.Branch
	if err := gitengine.Push(r.gitRunner, input.WorkspacePath, policy.Remote, policy.Branch); err != nil {
		msg := "git push failed"
		extra["message"] = msg
		return extra, core.OperationFailed, msg
	}
	extra["message"] = completionMessage
	return extra, core.OperationSucceeded, completionMessage
}

func toTerminalOperationState(state AgentRunState) (core.OperationState, error) {
	switch state {
	case AgentRunStateSucceeded:
		return core.OperationSucceeded, nil
	case AgentRunStateFailed:
		return core.OperationFailed, nil
	case AgentRunStateCancelled:
		return core.OperationCancelled, nil
	default:
		return "", fmt.Errorf("invalid agent run state %q", state)
	}
}