LocalAI/core/config/distributed_config_test.go

package config_test

import (
	"time"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"

	"github.com/mudler/LocalAI/core/config"
)

var _ = Describe("DistributedConfig backend NATS timeouts", func() {
	Context("BackendInstallTimeoutOrDefault", func() {
		It("returns 15 minutes when unset", func() {
			c := config.DistributedConfig{}
			Expect(c.BackendInstallTimeoutOrDefault()).To(Equal(15 * time.Minute))
		})

		It("returns the configured value when set", func() {
			c := config.DistributedConfig{BackendInstallTimeout: 42 * time.Minute}
			Expect(c.BackendInstallTimeoutOrDefault()).To(Equal(42 * time.Minute))
		})
	})

	Context("BackendUpgradeTimeoutOrDefault", func() {
		It("returns 15 minutes when unset", func() {
			c := config.DistributedConfig{}
			Expect(c.BackendUpgradeTimeoutOrDefault()).To(Equal(15 * time.Minute))
		})

		It("returns the configured value when set", func() {
			c := config.DistributedConfig{BackendUpgradeTimeout: 30 * time.Minute}
			Expect(c.BackendUpgradeTimeoutOrDefault()).To(Equal(30 * time.Minute))
		})
	})

	Context("ModelLoadTimeoutOrDefault", func() {
		It("returns 5 minutes when unset so existing clusters keep today's behaviour", func() {
			c := config.DistributedConfig{}
			Expect(c.ModelLoadTimeoutOrDefault()).To(Equal(5 * time.Minute))
		})

		It("returns the configured value when set", func() {
			c := config.DistributedConfig{ModelLoadTimeout: 45 * time.Minute}
			Expect(c.ModelLoadTimeoutOrDefault()).To(Equal(45 * time.Minute))
		})
	})
})

// Heartbeat checkpointing makes last_heartbeat up to one checkpoint interval
// stale by design, which is why the threshold defaults to 5 minutes. An
// operator who widens the checkpoint has to widen this to match, so it has to
// be reachable from the CLI rather than being a compile-time constant.
var _ = Describe("DistributedConfig stale node threshold", func() {
	It("defaults to 5 minutes, wide enough to cover a suppressed beat", func() {
		Expect(config.DistributedConfig{}.StaleNodeThresholdOrDefault()).
			To(Equal(5 * time.Minute))
		Expect(config.DefaultStaleNodeThreshold).
			To(BeNumerically(">", config.DefaultNodeHeartbeatCheckpoint),
				"a threshold at or below the checkpoint interval marks healthy, "+
					"beating nodes offline every cycle")
	})

	It("is configurable, so a widened checkpoint can be matched", func() {
		o := config.NewApplicationConfig(config.WithStaleNodeThreshold(20 * time.Minute))
		Expect(o.Distributed.StaleNodeThreshold).To(Equal(20 * time.Minute))
		Expect(o.Distributed.StaleNodeThresholdOrDefault()).To(Equal(20 * time.Minute))
	})
})

var _ = Describe("DistributedConfig flag-name constants", func() {
	// Pin the kebab-case strings so a rename of the Go field name (or a
	// CLI flag naming convention change) forces the constant to update,
	// keeping the Validate error messages and any future operator-facing
	// surface in sync with the actual CLI flag.
	DescribeTable("flag name constants",
		func(actual, expected string) {
			Expect(actual).To(Equal(expected))
		},
		Entry("MCP tool timeout", config.FlagMCPToolTimeout, "mcp-tool-timeout"),
		Entry("MCP discovery timeout", config.FlagMCPDiscoveryTimeout, "mcp-discovery-timeout"),
		Entry("worker wait timeout", config.FlagWorkerWaitTimeout, "worker-wait-timeout"),
		Entry("drain timeout", config.FlagDrainTimeout, "drain-timeout"),
		Entry("health check interval", config.FlagHealthCheckInterval, "health-check-interval"),
		Entry("stale node threshold", config.FlagStaleNodeThreshold, "stale-node-threshold"),
		Entry("node heartbeat checkpoint", config.FlagNodeHeartbeatCheckpoint, "node-heartbeat-checkpoint"),
		Entry("MCP CI job timeout", config.FlagMCPCIJobTimeout, "mcp-ci-job-timeout"),
		Entry("backend install timeout", config.FlagBackendInstallTimeout, "backend-install-timeout"),
		Entry("backend upgrade timeout", config.FlagBackendUpgradeTimeout, "backend-upgrade-timeout"),
		Entry("model load timeout", config.FlagModelLoadTimeout, "model-load-timeout"),
	)
})

var _ = Describe("DistributedConfig.Validate negative-duration errors", func() {
	It("rejects a negative BackendInstallTimeout with the flag name in the error", func() {
		c := config.DistributedConfig{
			Enabled:               true,
			NatsURL:               "nats://localhost:4222",
			BackendInstallTimeout: -1 * time.Second,
		}
		err := c.Validate()
		Expect(err).To(HaveOccurred())
		Expect(err.Error()).To(ContainSubstring(config.FlagBackendInstallTimeout))
		Expect(err.Error()).To(ContainSubstring("must not be negative"))
	})

	It("rejects a negative BackendUpgradeTimeout with the flag name in the error", func() {
		c := config.DistributedConfig{
			Enabled:               true,
			NatsURL:               "nats://localhost:4222",
			BackendUpgradeTimeout: -1 * time.Second,
		}
		err := c.Validate()
		Expect(err).To(HaveOccurred())
		Expect(err.Error()).To(ContainSubstring(config.FlagBackendUpgradeTimeout))
	})

	It("rejects a negative ModelLoadTimeout with the flag name in the error", func() {
		c := config.DistributedConfig{
			Enabled:          true,
			NatsURL:          "nats://localhost:4222",
			ModelLoadTimeout: -1 * time.Second,
		}
		err := c.Validate()
		Expect(err).To(HaveOccurred())
		Expect(err.Error()).To(ContainSubstring(config.FlagModelLoadTimeout))
		Expect(err.Error()).To(ContainSubstring("must not be negative"))
	})

	It("accepts all-zero durations as valid (defaults apply)", func() {
		c := config.DistributedConfig{
			Enabled: true,
			NatsURL: "nats://localhost:4222",
		}
		Expect(c.Validate()).To(Succeed())
	})
})

var _ = Describe("DistributedConfig.Validate registration auth", func() {
	It("rejects an empty registration token when RequireAuth is set", func() {
		c := config.DistributedConfig{
			Enabled:                 true,
			NatsURL:                 "nats://localhost:4222",
			RegistrationRequireAuth: true,
		}
		err := c.Validate()
		Expect(err).To(HaveOccurred())
		Expect(err.Error()).To(ContainSubstring("LOCALAI_REGISTRATION_REQUIRE_AUTH"))
		Expect(err.Error()).To(ContainSubstring("LOCALAI_REGISTRATION_TOKEN"))
	})

	It("accepts a set registration token when RequireAuth is set", func() {
		c := config.DistributedConfig{
			Enabled:                 true,
			NatsURL:                 "nats://localhost:4222",
			RegistrationToken:       "s3cret",
			RegistrationRequireAuth: true,
		}
		Expect(c.Validate()).To(Succeed())
	})

	It("warns but succeeds with an empty token when RequireAuth is unset", func() {
		c := config.DistributedConfig{
			Enabled: true,
			NatsURL: "nats://localhost:4222",
		}
		Expect(c.Validate()).To(Succeed())
	})

	It("rejects an empty token when the umbrella RequireAuth is set", func() {
		c := config.DistributedConfig{
			Enabled:     true,
			NatsURL:     "nats://localhost:4222",
			RequireAuth: true,
			// Provide NATS creds so only the registration-token gap remains.
			NatsServiceJWT:  "jwt",
			NatsServiceSeed: "seed",
			NatsAccountSeed: "acct",
		}
		err := c.Validate()
		Expect(err).To(HaveOccurred())
		Expect(err.Error()).To(ContainSubstring("LOCALAI_DISTRIBUTED_REQUIRE_AUTH"))
		Expect(err.Error()).To(ContainSubstring("LOCALAI_REGISTRATION_TOKEN"))
	})

	It("the umbrella implies NATS auth is required", func() {
		c := config.DistributedConfig{
			Enabled:           true,
			NatsURL:           "nats://localhost:4222",
			RegistrationToken: "tok", // registration layer satisfied
			RequireAuth:       true,  // umbrella → NATS creds now required
		}
		Expect(c.NatsAuthRequired()).To(BeTrue())
		Expect(c.RegistrationAuthRequired()).To(BeTrue())
		// Missing NATS service JWT/seed must now be fatal.
		err := c.Validate()
		Expect(err).To(HaveOccurred())
		Expect(err.Error()).To(ContainSubstring("LOCALAI_NATS_REQUIRE_AUTH"))
	})
})

var _ = Describe("DistributedConfig worker reconnect grace", func() {
	It("defaults clear of two ceiling backoffs plus the dial between them", func() {
		// The worker's own numbers (core/services/worker/tunnel.go): a 30s
		// backoff ceiling and a 10s dial budget, so two ceiling waits with a
		// hung dial between them puts the worker back at 70s. 60s would sit
		// under that and condemn a worker reconnecting exactly as designed;
		// 90s clears it with margin.
		Expect(config.DistributedConfig{}.ReconnectGraceOrDefault()).To(Equal(90 * time.Second))
		Expect(config.DefaultWorkerReconnectGrace).To(BeNumerically(">", 70*time.Second),
			"the default must clear two ceiling backoffs plus one handshake timeout")
	})

	It("takes a configured worker reconnect grace verbatim", func() {
		cfg := config.DistributedConfig{WorkerReconnectGrace: 5 * time.Minute}
		Expect(cfg.ReconnectGraceOrDefault()).To(Equal(5 * time.Minute))
	})

	It("falls back to the default rather than condemning every worker on a negative value", func() {
		// A negative grace makes every departure older than the window the
		// instant it is stamped, which reports a worker that has been gone for
		// two seconds as GONE, and gone is the one value a caller may reap on.
		cfg := config.DistributedConfig{WorkerReconnectGrace: -1 * time.Second}
		Expect(cfg.ReconnectGraceOrDefault()).To(Equal(config.DefaultWorkerReconnectGrace))
	})

	It("refuses to start on a negative grace rather than reaping on it", func() {
		// The flag is in Validate's negative-duration table, and this is what
		// says so. A negative grace makes every departure older than the window
		// the instant it is stamped, so the deployment would answer GONE for
		// every worker that has ever lost a tunnel, and gone is the one answer
		// a caller may reap and evict on.
		c := config.DistributedConfig{
			Enabled:              true,
			NatsURL:              "nats://localhost:4222",
			RegistrationToken:    "tok",
			WorkerReconnectGrace: -1 * time.Second,
		}
		err := c.Validate()
		Expect(err).To(HaveOccurred())
		Expect(err.Error()).To(ContainSubstring(config.FlagWorkerReconnectGrace))
	})

	It("is settable through the application option", func() {
		o := &config.ApplicationConfig{}
		config.WithWorkerReconnectGrace(90 * time.Second)(o)
		Expect(o.Distributed.ReconnectGraceOrDefault()).To(Equal(90 * time.Second))
	})
})