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)
}

// AgentRunPicker picks only agent_run-typed operations from the queue,
// preventing non-agent operations from being stranded in running state.
type AgentRunPicker interface {
	PickQueuedAgentRunOperation(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
	// BeforeRevision and AfterRevision, when both non-empty, let the scanner
	// emit a branch.updated event without requiring an existing cursor. This
	// covers the first push of a watched branch immediately after an agent_run.
	BeforeRevision string
	AfterRevision  string
}

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
	// AgentRunPicker, when set, is preferred over the generic OperationPicker so
	// that only agent_run operations are transitioned to running. If nil, the
	// caller-supplied generic picker is used instead.
	AgentRunPicker AgentRunPicker
}

type Runner struct {
	cfg                 config.Config
	logger              *slog.Logger
	picker              OperationPicker
	agentRunPicker      AgentRunPicker
	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,
		agentRunPicker:      deps.AgentRunPicker,
		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.agentRunPicker == nil && r.picker == nil {
		r.logger.Warn("worker picker not configured")
		return nil
	}

	var op core.Operation
	var ok bool
	var err error
	if r.agentRunPicker != nil {
		op, ok, err = r.agentRunPicker.PickQueuedAgentRunOperation(ctx, r.clock.Now())
	} else {
		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, loadErr := r.agentRunInputLoader.LoadAgentRunInput(ctx, op)
	if loadErr != nil {
		r.logger.Warn("agent_run input load failed; failing operation", "operation_id", op.ID)
		failAt := r.clock.Now().UTC()
		failedOp, ferr := r.operationStore.FailOperation(ctx, op.ID, failAt)
		if ferr != nil {
			return fmt.Errorf("fail agent_run operation after input load failure: %w", ferr)
		}
		failMsg := "agent run input load failed"
		failExtra := map[string]any{"message": failMsg}
		for _, evType := range []string{events.OperationFailed, events.AgentRunCompleted} {
			payload, perr := agentRunEventPayload(failedOp, core.OperationFailed, failAt, failExtra)
			if perr != nil {
				return perr
			}
			ev := storage.OperationEvent{
				OperationID: op.ID,
				Event: events.Event{
					ID:        eventID(op.ID, evType, failAt),
					Type:      evType,
					Subject:   operationEventSubject(op.ID),
					Payload:   payload,
					CreatedAt: failAt,
				},
			}
			if aerr := r.operationEventStore.AppendEvent(ctx, ev); aerr != nil {
				return fmt.Errorf("append %s event after input load failure: %w", evType, aerr)
			}
		}
		return nil
	}

	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(),
					BeforeRevision: evidence.BeforeRevision,
					AfterRevision:  evidence.AfterRevision,
				}
				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)
	}
}