gito/services/core/cmd/worker/main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"strings"

	"git.toki-labs.com/toki/gito/services/core/internal/config"
	"git.toki-labs.com/toki/gito/services/core/internal/controlplane"
	"git.toki-labs.com/toki/gito/services/core/internal/core"
	"git.toki-labs.com/toki/gito/services/core/internal/storage"
	"git.toki-labs.com/toki/gito/services/core/internal/worker"
)

// migrationSQL loads the initial migration SQL from well-known paths relative
// to the executable or common working directories, so the server/worker works
// regardless of the cwd it is launched from.
func migrationSQL() string {
	exe, err := os.Executable()
	if err != nil {
		exe = ""
	}
	candidates := []string{
		// set explicitly via env
		os.Getenv("GITO_MIGRATION_PATH"),
		// next to the binary (bin/gito-server → bin/../migrations/...)
		filepath.Join(filepath.Dir(exe), "migrations", "00001_initial.sql"),
		// launched from repo root
		filepath.Join("services", "core", "migrations", "00001_initial.sql"),
		// launched from services/core
		filepath.Join("migrations", "00001_initial.sql"),
	}
	for _, p := range candidates {
		if p == "" {
			continue
		}
		data, err := os.ReadFile(p)
		if err == nil {
			return string(data)
		}
	}
	return ""
}

// pgAgentRunInputLoader loads durable agent run inputs from the Postgres store
// and converts them to the worker input type.
type pgAgentRunInputLoader struct {
	store storage.AgentRunInputStore
}

func (l *pgAgentRunInputLoader) LoadAgentRunInput(ctx context.Context, op core.Operation) (worker.AgentRunInput, error) {
	durable, err := l.store.GetAgentRunInput(ctx, op.ID)
	if err != nil {
		if errors.Is(err, storage.ErrAgentRunInputNotFound) {
			return worker.AgentRunInput{}, fmt.Errorf("agent run input not found for operation %s: %w", op.ID, err)
		}
		return worker.AgentRunInput{}, fmt.Errorf("load agent run input: %w", err)
	}
	if strings.TrimSpace(durable.WorkspacePath) == "" || strings.TrimSpace(durable.Branch) == "" || strings.TrimSpace(durable.Instruction) == "" || strings.TrimSpace(durable.ExpectedRevision) == "" {
		return worker.AgentRunInput{}, fmt.Errorf("agent run input for %s is missing required fields (workspace_path, branch, instruction, expected_revision)", op.ID)
	}
	return worker.AgentRunInput{
		OperationID:      durable.OperationID,
		RepoID:           durable.RepoID,
		Branch:           durable.Branch,
		WorkspacePath:    durable.WorkspacePath,
		Instruction:      durable.Instruction,
		PolicyContext:    durable.PolicyContext,
		ExpectedRevision: durable.ExpectedRevision,
	}, nil
}

func main() {
	cfg := config.Load()
	logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))

	dsn := os.Getenv("DATABASE_URL")
	if dsn == "" {
		logger.Error("DATABASE_URL is required")
		os.Exit(1)
	}

	pgStore, err := storage.NewPgStore(context.Background(), dsn, migrationSQL())
	if err != nil {
		logger.Error("failed to open postgres store", "error", err)
		os.Exit(1)
	}
	logger.Info("postgres store connected")
	defer pgStore.Close()

	publisher := controlplane.NewOutboxBroadcaster(pgStore.OperationEvents())
	scanner := controlplane.NewRuntimeWithStore(publisher, pgStore)

	iopEndpoint := strings.TrimSpace(cfg.IOPEndpoint)
	if iopEndpoint == "" {
		logger.Info("IOP_ENDPOINT not configured; agent_run picking disabled")
		// Run with nil picker so worker logs the disabled message and exits cleanly.
		w := worker.NewRunner(cfg, logger, nil, nil, scanner)
		if err := w.Run(); err != nil {
			logger.Error("worker stopped", "error", err)
			os.Exit(1)
		}
		return
	}

	inputLoader := &pgAgentRunInputLoader{store: pgStore.AgentRunInputs()}
	invoker := worker.NewHTTPInvoker(iopEndpoint)

	w := worker.NewRunnerWithDependencies(cfg, logger, nil, nil, worker.RunnerDependencies{
		OperationStore:      pgStore.Operations(),
		OperationEventStore: pgStore.OperationEvents(),
		AgentRunPicker:      pgStore.Operations(),
		AgentRunInputLoader: inputLoader,
		AgentRunInvoker:     invoker,
		RevisionScanner:     scanner,
	})
	if err := w.Run(); err != nil {
		logger.Error("worker stopped", "error", err)
		os.Exit(1)
	}
}