serve.go

package cli

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"github.com/dvflw/mantle/internal/audit"
	"github.com/dvflw/mantle/internal/auth"
	"github.com/dvflw/mantle/internal/budget"
	"github.com/dvflw/mantle/internal/config"
	"github.com/dvflw/mantle/internal/connector"
	"github.com/dvflw/mantle/internal/db"
	"github.com/dvflw/mantle/internal/engine"
	"github.com/dvflw/mantle/internal/logging"
	"github.com/dvflw/mantle/internal/repo"
	reposync "github.com/dvflw/mantle/internal/repo/sync"
	"github.com/dvflw/mantle/internal/secret"
	"github.com/dvflw/mantle/internal/server"
	"github.com/spf13/cobra"
)

func newServeCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "serve",
		Short: "Start the Mantle server",
		Long:  "Starts a persistent process with the API server, cron scheduler, and webhook listener.",
		RunE: func(cmd *cobra.Command, args []string) error {
			cfg := config.FromContext(cmd.Context())
			if cfg == nil {
				return fmt.Errorf("config not loaded")
			}

			// Configure structured JSON logging before any other work.
			logging.Setup(cfg.Log.Level)

			database, err := db.Open(cfg.Database)
			if err != nil {
				return fmt.Errorf("failed to connect to database: %w", err)
			}
			defer database.Close()

			// Run migrations on startup.
			if err := db.Migrate(cmd.Context(), database); err != nil {
				return fmt.Errorf("running migrations: %w", err)
			}

			eng, err := engine.New(database)
			if err != nil {
				return fmt.Errorf("creating engine: %w", err)
			}
			eng.CEL.SetConfigEnv(cfg.Env)
			eng.MaxToolRoundsLimit = cfg.Engine.MaxToolRoundsLimit
			eng.MaxWorkflowDepth = cfg.Engine.MaxWorkflowDepth
			eng.MaxConcurrentExecutionsPerTeam = cfg.Engine.MaxConcurrentExecutionsPerTeam
			eng.RegisterWorkflowConnector()

			// Configure AWS defaults for AI (Bedrock) and S3 connectors.
			if aiConn, err := eng.Registry.Get("ai/completion"); err == nil {
				if ai, ok := aiConn.(*connector.AIConnector); ok {
					if cfg.AWS.Region != "" {
						ai.DefaultRegion = cfg.AWS.Region
						ai.AWSConfigFunc = connector.NewAWSConfig
					}
					ai.AllowedBaseURLs = cfg.Engine.AllowedBaseURLs
					ai.AllowedModels = cfg.Engine.AllowedModels
				}
			}

			// Apply the base-URL allowlist and AWS defaults to the embeddings
			// connector too, so ai/embed can't be pointed at arbitrary hosts and
			// its Bedrock provider gets a region. The model allowlist is
			// intentionally not shared: allowed_models lists chat models, and
			// applying it here would block embedding models.
			if embConn, err := eng.Registry.Get("ai/embed"); err == nil {
				if emb, ok := embConn.(*connector.EmbeddingConnector); ok {
					if cfg.AWS.Region != "" {
						emb.DefaultRegion = cfg.AWS.Region
						emb.AWSConfigFunc = connector.NewAWSConfig
					}
					emb.AllowedBaseURLs = cfg.Engine.AllowedBaseURLs
				}
			}

			// Wire up Postgres-backed audit emitter with auth context enrichment.
			auditor := &audit.PostgresEmitter{
				DB:                  database,
				AuthMethodExtractor: auth.AuthMethodFromContext,
			}
			eng.Auditor = auditor

			// Wire up credential resolver if encryption key is configured.
			if cfg.Encryption.Key != "" {
				encryptor, err := secret.NewEncryptor(cfg.Encryption.Key)
				if err != nil {
					return fmt.Errorf("configuring encryption: %w", err)
				}
				eng.Resolver = &secret.Resolver{
					Store: &secret.Store{DB: database, Encryptor: encryptor},
				}
			}

			// Wire budget enforcement into the engine.
			budgetStore := budget.NewStore(database)
			eng.BudgetStore = budgetStore
			eng.BudgetChecker = &budget.Checker{
				GlobalMonthlyTokenLimit:      cfg.Engine.Budget.GlobalMonthlyTokenLimit,
				DefaultTeamMonthlyTokenLimit: cfg.Engine.Budget.DefaultTeamMonthlyTokenLimit,
				ResetMode:                    cfg.Engine.Budget.ResetMode,
				ResetDay:                     cfg.Engine.Budget.ResetDay,
				GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
					u, err := budgetStore.GetGlobalTotalUsage(ctx, period)
					if err != nil {
						return 0, err
					}
					return u.TotalTokens, nil
				},
				GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) {
					u, err := budgetStore.GetUsage(ctx, teamID, provider, period)
					if err != nil {
						return 0, err
					}
					return u.TotalTokens, nil
				},
				GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) {
					return budgetStore.GetTeamBudget(ctx, teamID, provider)
				},
			}

			srv := server.New(database, eng, cfg.API.Address)
			srv.Auditor = auditor
			srv.TLSCertFile = cfg.API.TLS.CertFile
			srv.TLSKeyFile = cfg.API.TLS.KeyFile
			srv.AuthStore = &auth.Store{DB: database}
			srv.BudgetStore = budgetStore

			// Wire up OIDC validator if configured.
			if cfg.Auth.OIDC.IssuerURL != "" {
				oidcValidator, err := auth.NewOIDCValidator(
					cmd.Context(),
					cfg.Auth.OIDC.IssuerURL,
					cfg.Auth.OIDC.ClientID,
					cfg.Auth.OIDC.Audience,
					cfg.Auth.OIDC.AllowedDomains,
				)
				if err != nil {
					return fmt.Errorf("configuring OIDC: %w", err)
				}
				srv.OIDCValidator = oidcValidator
				fmt.Fprintf(cmd.OutOrStdout(), "OIDC authentication enabled (issuer: %s)\n", cfg.Auth.OIDC.IssuerURL)
			}

			// Handle SIGTERM and SIGINT for graceful shutdown.
			ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGTERM, syscall.SIGINT)
			defer stop()

			// Reconcile git_sync repos from config into the DB. Errors are fatal
			// at startup — a config mismatch must be fixed before the server runs.
			repoStore := &repo.Store{DB: database, Actor: "server"}
			if err := reposync.Reconcile(ctx, repoStore, cfg.GitSync.Repos); err != nil {
				return fmt.Errorf("reconciling git_sync repos: %w", err)
			}

			// Launch the git sync poller. It exits cleanly when ctx is cancelled.
			artifactBase := cfg.Storage.Path
			if artifactBase == "" {
				artifactBase = filepath.Join(os.TempDir(), "mantle-artifacts")
			}
			var secretResolver *secret.Resolver
			if cfg.Encryption.Key != "" {
				encryptor, encErr := secret.NewEncryptor(cfg.Encryption.Key)
				if encErr != nil {
					return fmt.Errorf("configuring encryption for git sync: %w", encErr)
				}
				secretResolver = &secret.Resolver{
					Store: &secret.Store{DB: database, Encryptor: encryptor},
				}
			}
			gitDriver := &reposync.GoGitDriver{
				BasePath: filepath.Join(artifactBase, "git"),
				Auth:     reposync.NewAuthResolver(ctx, secretResolver),
			}
			poller := &reposync.Poller{
				DB:     database,
				Store:  repoStore,
				Driver: gitDriver,
			}
			go poller.Run(ctx)

			// Wire git webhook handler into the server.
			srv.RepoStore = repoStore
			srv.GitDriver = gitDriver

			fmt.Fprintf(cmd.OutOrStdout(), "Mantle server starting on %s\n", cfg.API.Address)
			return srv.Start(ctx)
		},
	}
}