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

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
}

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

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

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

	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 = agentRunGitEvidencePayload(evidence)
			completionExtra["message"] = completionMessage
		}
	}

	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
}

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