refactoring and consolidation

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

Author: mudler

core/application/distributed.go

@@ -32,6 +32,7 @@ type DistributedServices struct {
 	FileMgr      *storage.FileManager
 	FileStager   nodes.FileStager
 	ModelAdapter *nodes.ModelRouterAdapter
+	Unloader     *nodes.RemoteUnloaderAdapter
 }
 
 // initDistributed validates distributed mode prerequisites and initializes
@@ -118,13 +119,17 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*Distribut
 	}
 	xlog.Info("Node registry initialized")
 
-	router := nodes.NewSmartRouter(registry)
+	// Collect SmartRouter option values; the router itself is created after all
+	// dependencies (including FileStager and Unloader) are ready.
+	var routerAuthToken string
 	if cfg.Distributed.RegistrationToken != "" {
-		router.SetAuthToken(cfg.Distributed.RegistrationToken)
+		routerAuthToken = cfg.Distributed.RegistrationToken
 	}
+	var routerGalleriesJSON string
 	if galleriesJSON, err := json.Marshal(cfg.BackendGalleries); err == nil {
-		router.SetGalleriesJSON(string(galleriesJSON))
+		routerGalleriesJSON = string(galleriesJSON)
 	}
+
 	healthMon := nodes.NewHealthMonitor(registry, authDB,
 		cfg.Distributed.HealthCheckIntervalOrDefault(),
 		cfg.Distributed.StaleNodeThresholdOrDefault(),
@@ -190,7 +195,17 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*Distribut
 		}, cfg.Distributed.RegistrationToken)
 		xlog.Info("File stager initialized (HTTP direct transfer)")
 	}
-	router.SetFileStager(fileStager)
+	// Create RemoteUnloaderAdapter — needed by SmartRouter and startup.go
+	remoteUnloader := nodes.NewRemoteUnloaderAdapter(registry, natsClient)
+
+	// All dependencies ready — build SmartRouter with all options at once
+	router := nodes.NewSmartRouter(registry, nodes.SmartRouterOptions{
+		Unloader:      remoteUnloader,
+		FileStager:    fileStager,
+		GalleriesJSON: routerGalleriesJSON,
+		AuthToken:     routerAuthToken,
+		DB:            authDB,
+	})
 
 	// Create ModelRouterAdapter to wire into ModelLoader
 	modelAdapter := nodes.NewModelRouterAdapter(router)
@@ -210,6 +225,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*Distribut
 		FileMgr:      fileMgr,
 		FileStager:   fileStager,
 		ModelAdapter: modelAdapter,
+		Unloader:     remoteUnloader,
 	}, nil
 }
 

core/application/startup.go

@@ -146,9 +146,7 @@ func New(opts ...config.AppOption) (*Application, error) {
 		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)
-		application.modelLoader.SetRemoteUnloader(remoteUnloader)
-		distSvc.Router.SetUnloader(remoteUnloader)
+		application.modelLoader.SetRemoteUnloader(distSvc.Unloader)
 		// Wire ModelRouter so grpcModel() delegates to SmartRouter in distributed mode
 		application.modelLoader.SetModelRouter(distSvc.ModelAdapter.AsModelRouter())
 		// Wire DistributedModelStore so shutdown/list/watchdog can find remote models
@@ -193,10 +191,10 @@ func New(opts ...config.AppOption) (*Application, error) {
 			}
 			// Wire distributed model/backend managers so delete propagates to workers
 			application.galleryService.SetModelManager(
-				nodes.NewDistributedModelManager(options, application.modelLoader, remoteUnloader),
+				nodes.NewDistributedModelManager(options, application.modelLoader, distSvc.Unloader),
 			)
 			application.galleryService.SetBackendManager(
-				nodes.NewDistributedBackendManager(options, application.modelLoader, remoteUnloader, distSvc.Registry),
+				nodes.NewDistributedBackendManager(options, application.modelLoader, distSvc.Unloader, distSvc.Registry),
 			)
 		}
 	}

core/cli/agent_worker.go

@@ -90,9 +90,7 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
 
 	// Start heartbeat
 	heartbeatInterval, _ := time.ParseDuration(cmd.HeartbeatInterval)
-	if heartbeatInterval == 0 {
-		heartbeatInterval = 10 * time.Second
-	}
+	heartbeatInterval = cmp.Or(heartbeatInterval, 10*time.Second)
 	// Context cancelled on shutdown — used by heartbeat and other background goroutines
 	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
 	defer shutdownCancel()
@@ -145,8 +143,8 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
 	// Subscribe to MCP CI job execution (load-balanced across agent workers).
 	// In distributed mode, MCP CI jobs are routed here because the frontend
 	// cannot create MCP sessions (e.g., stdio servers using docker).
-	natsClient.QueueSubscribe(messaging.SubjectJobsNew, messaging.QueueWorkers, func(data []byte) {
-		handleMCPCIJob(data, apiURL, cmd.APIToken, natsClient)
+	natsClient.QueueSubscribe(messaging.SubjectMCPCIJobsNew, messaging.QueueWorkers, func(data []byte) {
+		handleMCPCIJob(shutdownCtx, data, apiURL, cmd.APIToken, natsClient)
 	})
 
 	// Subscribe to backend stop events to clean up cached MCP sessions.
@@ -171,11 +169,7 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
 	xlog.Info("Shutting down agent worker")
 	dispatcher.Stop()
 	mcpTools.CloseAllMCPSessions()
-	if err := regClient.Deregister(context.Background(), nodeID); err != nil {
-		xlog.Warn("Failed to deregister", "error", err)
-	} else {
-		xlog.Info("Deregistered from frontend")
-	}
+	regClient.GracefulDeregister(nodeID)
 	return nil
 }
 
@@ -279,7 +273,9 @@ func sendMCPDiscoveryReply(reply func([]byte), servers []mcpRemote.MCPServerInfo
 	reply(data)
 }
 
-// natsClientAdapter wraps messaging.Client to satisfy NATSClient interface.
+// natsClientAdapter wraps messaging.Client to satisfy agents.NATSClient.
+// The concrete Client.QueueSubscribe returns *nats.Subscription; this adapter
+// widens it to messaging.Subscription so it matches the interface.
 type natsClientAdapter struct {
 	client *messaging.Client
 }
@@ -288,13 +284,13 @@ func (a *natsClientAdapter) Publish(subject string, data any) error {
 	return a.client.Publish(subject, data)
 }
 
-func (a *natsClientAdapter) QueueSubscribe(subject, queue string, handler func(data []byte)) (agents.NATSSub, error) {
+func (a *natsClientAdapter) QueueSubscribe(subject, queue string, handler func(data []byte)) (messaging.Subscription, error) {
 	return a.client.QueueSubscribe(subject, queue, handler)
 }
 
 // handleMCPCIJob processes an MCP CI job on the agent worker.
 // The agent worker can create MCP sessions (has docker) and call the LocalAI API for inference.
-func handleMCPCIJob(data []byte, apiURL, apiToken string, natsClient *messaging.Client) {
+func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken string, natsClient *messaging.Client) {
 	var evt jobs.JobEvent
 	if err := json.Unmarshal(data, &evt); err != nil {
 		xlog.Error("Failed to unmarshal job event", "error", err)
@@ -370,7 +366,7 @@ func handleMCPCIJob(data []byte, apiURL, apiToken string, natsClient *messaging.
 	llm := clients.NewLocalAILLM(task.Model, apiToken, apiURL)
 
 	// Build cogito options
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
+	ctx, cancel := context.WithTimeout(shutdownCtx, 10*time.Minute)
 	defer cancel()
 
 	// Update job status to running in DB

core/cli/worker.go

@@ -1,6 +1,7 @@
 package cli
 
 import (
+	"cmp"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -108,9 +109,7 @@ func (cmd *WorkerCMD) Run(ctx *cliContext.Context) error {
 
 	xlog.Info("Registered with frontend", "nodeID", nodeID, "frontend", cmd.RegisterTo)
 	heartbeatInterval, _ := time.ParseDuration(cmd.HeartbeatInterval)
-	if heartbeatInterval == 0 {
-		heartbeatInterval = 10 * time.Second
-	}
+	heartbeatInterval = cmp.Or(heartbeatInterval, 10*time.Second)
 	// Context cancelled on shutdown — used by heartbeat and other background goroutines
 	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
 	defer shutdownCancel()
@@ -121,7 +120,7 @@ func (cmd *WorkerCMD) Run(ctx *cliContext.Context) error {
 	httpAddr := cmd.resolveHTTPAddr()
 	stagingDir := filepath.Join(cmd.ModelsPath, "..", "staging")
 	dataDir := filepath.Join(cmd.ModelsPath, "..", "data")
-	httpServer, err := nodes.StartFileTransferServer(httpAddr, stagingDir, cmd.ModelsPath, dataDir, cmd.RegistrationToken, ml.BackendLogs())
+	httpServer, err := nodes.StartFileTransferServer(httpAddr, stagingDir, cmd.ModelsPath, dataDir, cmd.RegistrationToken, config.DefaultMaxUploadSize, ml.BackendLogs())
 	if err != nil {
 		return fmt.Errorf("starting HTTP file transfer server: %w", err)
 	}
@@ -177,9 +176,9 @@ func (cmd *WorkerCMD) Run(ctx *cliContext.Context) error {
 	<-sigCh
 
 	xlog.Info("Shutting down worker")
+	regClient.GracefulDeregister(nodeID)
 	supervisor.stopAllBackends()
 	nodes.ShutdownFileTransferServer(httpServer)
-	regClient.GracefulDeregister(nodeID)
 	return nil
 }
 

core/config/distributed_config.go

@@ -28,6 +28,8 @@ type DistributedConfig struct {
 	DrainTimeout        time.Duration // Time to wait for in-flight requests during drain (default 30s)
 	HealthCheckInterval time.Duration // Health monitor check interval (default 15s)
 	StaleNodeThreshold  time.Duration // Time before a node is considered stale (default 60s)
+
+	MaxUploadSize int64 // Maximum upload body size in bytes (default 50 GB)
 }
 
 // Distributed config options
@@ -98,6 +100,9 @@ const (
 	DefaultStaleNodeThreshold  = 60 * time.Second
 )
 
+// DefaultMaxUploadSize is the default maximum upload body size (50 GB).
+const DefaultMaxUploadSize int64 = 50 << 30
+
 // MCPToolTimeoutOrDefault returns the configured timeout or the default.
 func (c DistributedConfig) MCPToolTimeoutOrDefault() time.Duration {
 	return cmp.Or(c.MCPToolTimeout, DefaultMCPToolTimeout)
@@ -127,3 +132,11 @@ func (c DistributedConfig) HealthCheckIntervalOrDefault() time.Duration {
 func (c DistributedConfig) StaleNodeThresholdOrDefault() time.Duration {
 	return cmp.Or(c.StaleNodeThreshold, DefaultStaleNodeThreshold)
 }
+
+// MaxUploadSizeOrDefault returns the configured max upload size or the default.
+func (c DistributedConfig) MaxUploadSizeOrDefault() int64 {
+	if c.MaxUploadSize > 0 {
+		return c.MaxUploadSize
+	}
+	return DefaultMaxUploadSize
+}

core/config/model_config.go

@@ -105,6 +105,11 @@ type AgentConfig struct {
 	ForceReasoningTool    bool `yaml:"force_reasoning_tool,omitempty" json:"force_reasoning_tool,omitempty"`
 }
 
+// HasMCPServers returns true if any MCP servers (remote or stdio) are configured.
+func (c MCPConfig) HasMCPServers() bool {
+	return c.Servers != "" || c.Stdio != ""
+}
+
 func (c *MCPConfig) MCPConfigFromYAML() (MCPGenericConfig[MCPRemoteServers], MCPGenericConfig[MCPSTDIOServers], error) {
 	var remote MCPGenericConfig[MCPRemoteServers]
 	var stdio MCPGenericConfig[MCPSTDIOServers]

core/http/endpoints/localai/nodes.go

@@ -473,7 +473,18 @@ func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToke
 // /v1/backend-logs/{modelId}/ws endpoint for real-time log streaming.
 func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken string) echo.HandlerFunc {
 	browserUpgrader := websocket.Upgrader{
-		CheckOrigin: func(r *http.Request) bool { return true },
+		CheckOrigin: func(r *http.Request) bool {
+			origin := r.Header.Get("Origin")
+			if origin == "" {
+				return true // no origin header = same-origin or non-browser
+			}
+			// Parse origin URL and compare host with request host
+			u, err := url.Parse(origin)
+			if err != nil {
+				return false
+			}
+			return u.Host == r.Host
+		},
 	}
 
 	return func(c echo.Context) error {
@@ -507,14 +518,12 @@ func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken s
 			return nil
 		}
 
-		// Use sync.Once wrappers to avoid double-close and ensure each
+		// Use sync.OnceFunc wrappers to avoid double-close and ensure each
 		// goroutine can safely close the *other* connection to unblock
 		// its peer's ReadMessage call.
 		done := make(chan struct{})
-		var closeWorkerOnce sync.Once
-		closeWorker := func() { closeWorkerOnce.Do(func() { workerWS.Close() }) }
-		var closeBrowserOnce sync.Once
-		closeBrowser := func() { closeBrowserOnce.Do(func() { browserWS.Close() }) }
+		closeWorker := sync.OnceFunc(func() { workerWS.Close() })
+		closeBrowser := sync.OnceFunc(func() { browserWS.Close() })
 
 		// Worker → Browser
 		go func() {

core/http/routes/nodes.go

@@ -28,6 +28,11 @@ func nodeReadyMiddleware(registry *nodes.NodeRegistry) echo.MiddlewareFunc {
 // RegisterNodeSelfServiceRoutes registers /api/node/ endpoints used by backend
 // nodes themselves (register, heartbeat, drain, query own models, deregister).
 // These are authenticated via the registration token, not admin middleware.
+//
+// TODO(security): Self-service routes currently accept any node :id with a valid
+// registration token. A compromised worker can heartbeat/drain/deregister other nodes.
+// To fix: issue per-node bearer tokens during registration and validate :id matches
+// the authenticated node. See security review item S2.
 func RegisterNodeSelfServiceRoutes(e *echo.Echo, registry *nodes.NodeRegistry, registrationToken string, autoApprove bool, authDB *gorm.DB, hmacSecret string) {
 	if registry == nil {
 		return