refactoring and consolidation

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

Author: mudler

core/application/agent_jobs.go

@@ -30,11 +30,13 @@ func (a *Application) RestartAgentJobService() error {
 	)
 
 	// Re-apply distributed wiring if available (matches startup.go logic)
-	if a.jobDispatcher != nil {
-		agentJobService.SetDistributedBackends(a.jobDispatcher)
-	}
-	if a.jobStore != nil {
-		agentJobService.SetDistributedJobStore(a.jobStore)
+	if d := a.Distributed(); d != nil {
+		if d.Dispatcher != nil {
+			agentJobService.SetDistributedBackends(d.Dispatcher)
+		}
+		if d.JobStore != nil {
+			agentJobService.SetDistributedJobStore(d.JobStore)
+		}
 	}
 
 	// Start the service

core/application/application.go

@@ -2,20 +2,16 @@ package application
 
 import (
 	"context"
+	"math/rand/v2"
 	"sync"
 	"sync/atomic"
 	"time"
 
 	"github.com/mudler/LocalAI/core/config"
 	mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp"
-	"github.com/mudler/LocalAI/core/services/agents"
 	"github.com/mudler/LocalAI/core/services/agentpool"
 	"github.com/mudler/LocalAI/core/services/galleryop"
-	"github.com/mudler/LocalAI/core/services/distributed"
-	"github.com/mudler/LocalAI/core/services/jobs"
-	"github.com/mudler/LocalAI/core/services/messaging"
 	"github.com/mudler/LocalAI/core/services/nodes"
-	"github.com/mudler/LocalAI/core/services/storage"
 	"github.com/mudler/LocalAI/core/templates"
 	"github.com/mudler/LocalAI/pkg/model"
 	"github.com/mudler/xlog"
@@ -40,17 +36,7 @@ type Application struct {
 	agentJobMutex      sync.Mutex
 
 	// Distributed mode services (nil when not in distributed mode)
-	natsClient     *messaging.Client
-	objectStore    storage.ObjectStore
-	nodeRegistry   *nodes.NodeRegistry
-	smartRouter    *nodes.SmartRouter
-	healthMon      *nodes.HealthMonitor
-	jobStore       *jobs.JobStore
-	jobDispatcher  *jobs.Dispatcher
-	agentStore     *agents.AgentStore
-	agentBridge    *agents.EventBridge
-	distStores     *distributed.Stores
-	fileManager    *storage.FileManager
+	distributed *DistributedServices
 }
 
 func newApplication(appConfig *config.ApplicationConfig) *Application {
@@ -107,77 +93,27 @@ func (a *Application) StartupConfig() *config.ApplicationConfig {
 	return a.startupConfig
 }
 
-// NatsClient returns the NATS messaging client, or nil if not in distributed mode.
-func (a *Application) NatsClient() *messaging.Client {
-	return a.natsClient
-}
-
-// ObjectStore returns the object storage backend, or nil if not in distributed mode.
-func (a *Application) ObjectStore() storage.ObjectStore {
-	return a.objectStore
-}
-
-// NodeRegistry returns the node registry, or nil if not in distributed mode.
-func (a *Application) NodeRegistry() *nodes.NodeRegistry {
-	return a.nodeRegistry
-}
-
-// SmartRouter returns the smart router, or nil if not in distributed mode.
-func (a *Application) SmartRouter() *nodes.SmartRouter {
-	return a.smartRouter
-}
-
-// JobStore returns the distributed job store, or nil if not in distributed mode.
-func (a *Application) JobStore() *jobs.JobStore {
-	return a.jobStore
-}
-
-// JobDispatcher returns the distributed job dispatcher, or nil if not in distributed mode.
-func (a *Application) JobDispatcher() *jobs.Dispatcher {
-	return a.jobDispatcher
-}
-
-// AgentStore returns the distributed agent store, or nil if not in distributed mode.
-func (a *Application) AgentStore() *agents.AgentStore {
-	return a.agentStore
-}
-
-// AgentEventBridge returns the agent event bridge, or nil if not in distributed mode.
-func (a *Application) AgentEventBridge() *agents.EventBridge {
-	return a.agentBridge
-}
-
-// DistributedStores returns the Phase 4 distributed stores, or nil if not in distributed mode.
-func (a *Application) DistributedStores() *distributed.Stores {
-	return a.distStores
-}
-
-// FileManager returns the file manager for object storage, or nil if not in distributed mode.
-func (a *Application) FileManager() *storage.FileManager {
-	return a.fileManager
-}
-
-// HealthMonitor returns the health monitor, or nil if not in distributed mode.
-func (a *Application) HealthMonitor() *nodes.HealthMonitor {
-	return a.healthMon
+// Distributed returns the distributed services, or nil if not in distributed mode.
+func (a *Application) Distributed() *DistributedServices {
+	return a.distributed
 }
 
 // IsDistributed returns true if the application is running in distributed mode.
 func (a *Application) IsDistributed() bool {
-	return a.applicationConfig.Distributed.Enabled
+	return a.distributed != nil
 }
 
 // waitForHealthyWorker blocks until at least one healthy backend worker is registered.
 // This prevents the agent pool from failing during startup when workers haven't connected yet.
 func (a *Application) waitForHealthyWorker() {
 	maxWait := a.applicationConfig.Distributed.WorkerWaitTimeoutOrDefault()
-	const pollInterval = 2 * time.Second
+	const basePoll = 2 * time.Second
 
 	xlog.Info("Waiting for at least one healthy backend worker before starting agent pool")
 	deadline := time.Now().Add(maxWait)
 
 	for time.Now().Before(deadline) {
-		registered, err := a.nodeRegistry.List()
+		registered, err := a.distributed.Registry.List()
 		if err == nil {
 			for _, n := range registered {
 				if n.NodeType == nodes.NodeTypeBackend && n.Status == nodes.StatusHealthy {
@@ -186,10 +122,12 @@ func (a *Application) waitForHealthyWorker() {
 				}
 			}
 		}
+		// Add 0-1s jitter to prevent thundering-herd on the node registry
+		jitter := time.Duration(rand.Int64N(int64(time.Second)))
 		select {
 		case <-a.applicationConfig.Context.Done():
 			return
-		case <-time.After(pollInterval):
+		case <-time.After(basePoll + jitter):
 		}
 	}
 	xlog.Warn("No healthy backend worker found after waiting, proceeding anyway")
@@ -237,24 +175,23 @@ func (a *Application) StartAgentPool() {
 	if a.authDB != nil {
 		aps.SetAuthDB(a.authDB)
 	}
-	if a.distStores != nil && a.distStores.Skills != nil {
-		aps.SetSkillStore(a.distStores.Skills)
-	}
 	// Wire distributed mode components
-	if a.natsClient != nil {
-		aps.SetNATSClient(a.natsClient)
-	}
-	if a.agentBridge != nil {
-		aps.SetEventBridge(a.agentBridge)
-	}
-	if a.agentStore != nil {
-		aps.SetAgentStore(a.agentStore)
-	}
-	// In distributed mode, wait for at least one healthy backend worker before
-	// starting the agent pool. Collections initialization calls embeddings which
-	// require a worker to be registered and subscribed to NATS.
-	if a.IsDistributed() && a.nodeRegistry != nil {
-		a.waitForHealthyWorker()
+	if d := a.Distributed(); d != nil {
+		if d.DistStores != nil && d.DistStores.Skills != nil {
+			aps.SetSkillStore(d.DistStores.Skills)
+		}
+		aps.SetNATSClient(d.Nats)
+		if d.AgentBridge != nil {
+			aps.SetEventBridge(d.AgentBridge)
+		}
+		if d.AgentStore != nil {
+			aps.SetAgentStore(d.AgentStore)
+		}
+		// Wait for at least one healthy backend worker before starting the agent pool.
+		// Collections initialization calls embeddings which require a worker.
+		if d.Registry != nil {
+			a.waitForHealthyWorker()
+		}
 	}
 
 	if err := aps.Start(a.applicationConfig.Context); err != nil {

core/application/distributed.go

@@ -17,27 +17,27 @@ import (
 	"gorm.io/gorm"
 )
 
-// distributedServices holds all services initialized for distributed mode.
-type distributedServices struct {
-	nats         *messaging.Client
-	store        storage.ObjectStore
-	registry     *nodes.NodeRegistry
-	router       *nodes.SmartRouter
-	health       *nodes.HealthMonitor
-	jobStore     *jobs.JobStore
-	dispatcher   *jobs.Dispatcher
-	agentStore   *agents.AgentStore
-	agentBridge  *agents.EventBridge
-	distStores   *distributed.Stores
-	fileMgr      *storage.FileManager
-	fileStager   nodes.FileStager
-	modelAdapter *nodes.ModelRouterAdapter
+// DistributedServices holds all services initialized for distributed mode.
+type DistributedServices struct {
+	Nats         *messaging.Client
+	Store        storage.ObjectStore
+	Registry     *nodes.NodeRegistry
+	Router       *nodes.SmartRouter
+	Health       *nodes.HealthMonitor
+	JobStore     *jobs.JobStore
+	Dispatcher   *jobs.Dispatcher
+	AgentStore   *agents.AgentStore
+	AgentBridge  *agents.EventBridge
+	DistStores   *distributed.Stores
+	FileMgr      *storage.FileManager
+	FileStager   nodes.FileStager
+	ModelAdapter *nodes.ModelRouterAdapter
 }
 
 // initDistributed validates distributed mode prerequisites and initializes
 // NATS, object storage, node registry, and instance identity.
 // Returns nil if distributed mode is not enabled.
-func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*distributedServices, error) {
+func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*DistributedServices, error) {
 	if !cfg.Distributed.Enabled {
 		return nil, nil
 	}
@@ -196,20 +196,20 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*distribut
 	modelAdapter := nodes.NewModelRouterAdapter(router)
 
 	success = true
-	return &distributedServices{
-		nats:         natsClient,
-		store:        store,
-		registry:     registry,
-		router:       router,
-		health:       healthMon,
-		jobStore:     jobStore,
-		dispatcher:   dispatcher,
-		agentStore:   agentStore,
-		agentBridge:  agentBridge,
-		distStores:   distStores,
-		fileMgr:      fileMgr,
-		fileStager:   fileStager,
-		modelAdapter: modelAdapter,
+	return &DistributedServices{
+		Nats:         natsClient,
+		Store:        store,
+		Registry:     registry,
+		Router:       router,
+		Health:       healthMon,
+		JobStore:     jobStore,
+		Dispatcher:   dispatcher,
+		AgentStore:   agentStore,
+		AgentBridge:  agentBridge,
+		DistStores:   distStores,
+		FileMgr:      fileMgr,
+		FileStager:   fileStager,
+		ModelAdapter: modelAdapter,
 	}, nil
 }
 

core/application/startup.go

@@ -143,70 +143,60 @@ func New(opts ...config.AppOption) (*Application, error) {
 		return nil, fmt.Errorf("distributed mode initialization failed: %w", err)
 	}
 	if distSvc != nil {
-		application.natsClient = distSvc.nats
-		application.objectStore = distSvc.store
-		application.nodeRegistry = distSvc.registry
-		application.smartRouter = distSvc.router
-		application.healthMon = distSvc.health
-		application.jobStore = distSvc.jobStore
-		application.jobDispatcher = distSvc.dispatcher
-		application.agentStore = distSvc.agentStore
-		application.agentBridge = distSvc.agentBridge
-		application.distStores = distSvc.distStores
-		application.fileManager = distSvc.fileMgr
+		application.distributed = distSvc
 		// Wire remote model unloader so ShutdownModel works for remote nodes
 		// Uses NATS to tell serve-backend nodes to Free + kill their backend process
-		remoteUnloader := nodes.NewRemoteUnloaderAdapter(distSvc.registry, distSvc.nats)
+		remoteUnloader := nodes.NewRemoteUnloaderAdapter(distSvc.Registry, distSvc.Nats)
 		application.modelLoader.SetRemoteUnloader(remoteUnloader)
-		distSvc.router.SetUnloader(remoteUnloader)
+		distSvc.Router.SetUnloader(remoteUnloader)
 		// Wire ModelRouter so grpcModel() delegates to SmartRouter in distributed mode
-		application.modelLoader.SetModelRouter(distSvc.modelAdapter.AsModelRouter())
+		application.modelLoader.SetModelRouter(distSvc.ModelAdapter.AsModelRouter())
 		// Wire DistributedModelStore so shutdown/list/watchdog can find remote models
 		distStore := nodes.NewDistributedModelStore(
 			model.NewInMemoryModelStore(),
-			distSvc.registry,
+			distSvc.Registry,
 		)
 		application.modelLoader.SetModelStore(distStore)
 		// Start health monitor
-		distSvc.health.Start(options.Context)
+		distSvc.Health.Start(options.Context)
 		// In distributed mode, MCP CI jobs are executed by agent workers (not the frontend)
 		// because the frontend can't create MCP sessions (e.g., stdio servers using docker).
 		// The dispatcher still subscribes to jobs.new for persistence (result/progress subs)
 		// but does NOT set a workerFn — agent workers consume jobs from the same NATS queue.
 
 		// Wire model config loader so job events include model config for agent workers
-		distSvc.dispatcher.SetModelConfigLoader(application.backendLoader)
+		distSvc.Dispatcher.SetModelConfigLoader(application.backendLoader)
 
 		// Start job dispatcher — abort startup if it fails, as jobs would be accepted but never dispatched
-		if err := distSvc.dispatcher.Start(options.Context); err != nil {
+		if err := distSvc.Dispatcher.Start(options.Context); err != nil {
 			return nil, fmt.Errorf("starting job dispatcher: %w", err)
 		}
 		// Start ephemeral file cleanup
-		storage.StartEphemeralCleanup(options.Context, distSvc.fileMgr, 0, 0)
+		storage.StartEphemeralCleanup(options.Context, distSvc.FileMgr, 0, 0)
 		// Wire distributed backends into AgentJobService (before Start)
 		if application.agentJobService != nil {
-			application.agentJobService.SetDistributedBackends(distSvc.dispatcher)
-			application.agentJobService.SetDistributedJobStore(distSvc.jobStore)
+			application.agentJobService.SetDistributedBackends(distSvc.Dispatcher)
+			application.agentJobService.SetDistributedJobStore(distSvc.JobStore)
 		}
 		// Wire skill store into AgentPoolService (wired at pool start time via closure)
 		// The actual wiring happens in StartAgentPool since the pool doesn't exist yet.
 
 		// Wire NATS and gallery store into GalleryService for cross-instance progress/cancel
 		if application.galleryService != nil {
-			application.galleryService.SetNATSClient(distSvc.nats)
-			if distSvc.distStores != nil && distSvc.distStores.Gallery != nil {
+			application.galleryService.SetNATSClient(distSvc.Nats)
+			if distSvc.DistStores != nil && distSvc.DistStores.Gallery != nil {
 				// Clean up stale in-progress operations from previous crashed instances
-				if err := distSvc.distStores.Gallery.CleanStale(30 * time.Minute); err != nil {
+				if err := distSvc.DistStores.Gallery.CleanStale(30 * time.Minute); err != nil {
 					xlog.Warn("Failed to clean stale gallery operations", "error", err)
 				}
-				application.galleryService.SetGalleryStore(distSvc.distStores.Gallery)
+				application.galleryService.SetGalleryStore(distSvc.DistStores.Gallery)
 			}
 			// Wire distributed model/backend managers so delete propagates to workers
 			application.galleryService.SetModelManager(
 				nodes.NewDistributedModelManager(options, application.modelLoader, remoteUnloader),
 			)
 			application.galleryService.SetBackendManager(
-				nodes.NewDistributedBackendManager(options, application.modelLoader, remoteUnloader, distSvc.registry),
+				nodes.NewDistributedBackendManager(options, application.modelLoader, remoteUnloader, distSvc.Registry),
 			)
 		}
 	}