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

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

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()
	}
	return &Runner{
		cfg:                 cfg,
		logger:              logger,
		picker:              picker,
		clock:               clock,
		operationStore:      deps.OperationStore,
		operationEventStore: deps.OperationEventStore,
		agentRunInputLoader: deps.AgentRunInputLoader,
		agentRunInvoker:     deps.AgentRunInvoker,
	}
}

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

	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, map[string]any{"message": completionMessage})
	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, map[string]any{"message": completionMessage})
	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 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)
	}
}