refactoring and consolidation

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

Author: mudler

core/application/application.go

@@ -113,7 +113,7 @@ func (a *Application) waitForHealthyWorker() {
 	deadline := time.Now().Add(maxWait)
 
 	for time.Now().Before(deadline) {
-		registered, err := a.distributed.Registry.List()
+		registered, err := a.distributed.Registry.List(context.Background())
 		if err == nil {
 			for _, n := range registered {
 				if n.NodeType == nodes.NodeTypeBackend && n.Status == nodes.StatusHealthy {

core/application/distributed.go

@@ -1,6 +1,7 @@
 package application
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"strings"
@@ -36,6 +37,26 @@ type DistributedServices struct {
 	Unloader     *nodes.RemoteUnloaderAdapter
 }
 
+// Shutdown stops all distributed services in reverse initialization order.
+// It is safe to call on a nil receiver.
+func (ds *DistributedServices) Shutdown() {
+	if ds == nil {
+		return
+	}
+	if ds.Health != nil {
+		ds.Health.Stop()
+	}
+	if ds.Dispatcher != nil {
+		ds.Dispatcher.Stop()
+	}
+	// AgentBridge has no Close method — its NATS subscriptions are cleaned up
+	// when the NATS client is closed below.
+	if ds.Nats != nil {
+		ds.Nats.Close()
+	}
+	xlog.Info("Distributed services shut down")
+}
+
 // initDistributed validates distributed mode prerequisites and initializes
 // NATS, object storage, node registry, and instance identity.
 // Returns nil if distributed mode is not enabled.
@@ -46,6 +67,11 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*Distribut
 
 	xlog.Info("Distributed mode enabled — validating prerequisites")
 
+	// Validate distributed config (NATS URL, S3 credential pairing, durations, etc.)
+	if err := cfg.Distributed.Validate(); err != nil {
+		return nil, err
+	}
+
 	// Validate PostgreSQL is configured (auth DB must be PostgreSQL for distributed mode)
 	if !cfg.Auth.Enabled {
 		return nil, fmt.Errorf("distributed mode requires authentication to be enabled (--auth / LOCALAI_AUTH=true)")
@@ -54,11 +80,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*Distribut
 		return nil, fmt.Errorf("distributed mode requires PostgreSQL for auth database (got %q)", sanitize.URL(cfg.Auth.DatabaseURL))
 	}
 
-	// Validate NATS
-	if cfg.Distributed.NatsURL == "" {
-		return nil, fmt.Errorf("distributed mode requires --nats-url / LOCALAI_NATS_URL")
-	}
-
 	// Generate instance ID if not set
 	if cfg.Distributed.InstanceID == "" {
 		cfg.Distributed.InstanceID = uuid.New().String()
@@ -134,6 +155,8 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*Distribut
 	healthMon := nodes.NewHealthMonitor(registry, authDB,
 		cfg.Distributed.HealthCheckIntervalOrDefault(),
 		cfg.Distributed.StaleNodeThresholdOrDefault(),
+		routerAuthToken,
+		cfg.Distributed.PerModelHealthCheck,
 	)
 
 	// Initialize job store
@@ -185,7 +208,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*Distribut
 		xlog.Info("File stager initialized (S3+NATS)")
 	} else {
 		fileStager = nodes.NewHTTPFileStager(func(nodeID string) (string, error) {
-			node, err := registry.Get(nodeID)
+			node, err := registry.Get(context.Background(), nodeID)
 			if err != nil {
 				return "", err
 			}

core/application/startup.go

@@ -268,6 +268,7 @@ func New(opts ...config.AppOption) (*Application, error) {
 	go func() {
 		<-options.Context.Done()
 		xlog.Debug("Context canceled, shutting down")
+		application.distributed.Shutdown()
 		err := application.ModelLoader().StopAllGRPC()
 		if err != nil {
 			xlog.Error("error while stopping all grpc backends", "error", err)

core/cli/agent_worker.go

@@ -51,6 +51,9 @@ type AgentWorkerCMD struct {
 	// NATS subjects
 	Subject string `env:"LOCALAI_AGENT_SUBJECT" default:"agent.execute" help:"NATS subject for agent execution" group:"distributed"`
 	Queue   string `env:"LOCALAI_AGENT_QUEUE" default:"agent-workers" help:"NATS queue group name" group:"distributed"`
+
+	// Timeouts
+	MCPCIJobTimeout string `env:"LOCALAI_MCP_CI_JOB_TIMEOUT" default:"10m" help:"Timeout for MCP CI job execution" group:"distributed"`
 }
 
 func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
@@ -90,7 +93,10 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
 	}
 
 	// Start heartbeat
-	heartbeatInterval, _ := time.ParseDuration(cmd.HeartbeatInterval)
+	heartbeatInterval, err := time.ParseDuration(cmd.HeartbeatInterval)
+	if err != nil && cmd.HeartbeatInterval != "" {
+		xlog.Warn("invalid heartbeat interval, using default 10s", "input", cmd.HeartbeatInterval, "error", err)
+	}
 	heartbeatInterval = cmp.Or(heartbeatInterval, 10*time.Second)
 	// Context cancelled on shutdown — used by heartbeat and other background goroutines
 	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
@@ -145,8 +151,14 @@ 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).
+	mcpCIJobTimeout, err := time.ParseDuration(cmd.MCPCIJobTimeout)
+	if err != nil && cmd.MCPCIJobTimeout != "" {
+		xlog.Warn("invalid MCP CI job timeout, using default 10m", "input", cmd.MCPCIJobTimeout, "error", err)
+	}
+	mcpCIJobTimeout = cmp.Or(mcpCIJobTimeout, config.DefaultMCPCIJobTimeout)
+
 	natsClient.QueueSubscribe(messaging.SubjectMCPCIJobsNew, messaging.QueueWorkers, func(data []byte) {
-		handleMCPCIJob(shutdownCtx, data, apiURL, cmd.APIToken, natsClient)
+		handleMCPCIJob(shutdownCtx, data, apiURL, cmd.APIToken, natsClient, mcpCIJobTimeout)
 	})
 
 	// Subscribe to backend stop events to clean up cached MCP sessions.
@@ -277,7 +289,7 @@ func sendMCPDiscoveryReply(reply func([]byte), servers []mcpRemote.MCPServerInfo
 
 // 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(shutdownCtx context.Context, data []byte, apiURL, apiToken string, natsClient *messaging.Client) {
+func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken string, natsClient messaging.MessagingClient, jobTimeout time.Duration) {
 	var evt jobs.JobEvent
 	if err := json.Unmarshal(data, &evt); err != nil {
 		xlog.Error("Failed to unmarshal job event", "error", err)
@@ -353,7 +365,7 @@ func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken s
 	llm := clients.NewLocalAILLM(task.Model, apiToken, apiURL)
 
 	// Build cogito options
-	ctx, cancel := context.WithTimeout(shutdownCtx, 10*time.Minute)
+	ctx, cancel := context.WithTimeout(shutdownCtx, jobTimeout)
 	defer cancel()
 
 	// Update job status to running in DB
@@ -433,11 +445,11 @@ func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken s
 	xlog.Info("MCP CI job completed", "jobID", evt.JobID, "resultLen", len(result))
 }
 
-func publishJobStatus(nc *messaging.Client, jobID, status, message string) {
+func publishJobStatus(nc messaging.MessagingClient, jobID, status, message string) {
 	jobs.PublishJobProgress(nc, jobID, status, message)
 }
 
-func publishJobResult(nc *messaging.Client, jobID, status, result, errMsg string) {
+func publishJobResult(nc messaging.MessagingClient, jobID, status, result, errMsg string) {
 	jobs.PublishJobResult(nc, jobID, status, result, errMsg)
 }
 

core/cli/worker.go

@@ -109,7 +109,10 @@ func (cmd *WorkerCMD) Run(ctx *cliContext.Context) error {
 	}
 
 	xlog.Info("Registered with frontend", "nodeID", nodeID, "frontend", cmd.RegisterTo)
-	heartbeatInterval, _ := time.ParseDuration(cmd.HeartbeatInterval)
+	heartbeatInterval, err := time.ParseDuration(cmd.HeartbeatInterval)
+	if err != nil && cmd.HeartbeatInterval != "" {
+		xlog.Warn("invalid heartbeat interval, using default 10s", "input", cmd.HeartbeatInterval, "error", err)
+	}
 	heartbeatInterval = cmp.Or(heartbeatInterval, 10*time.Second)
 	// Context cancelled on shutdown — used by heartbeat and other background goroutines
 	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
@@ -184,7 +187,7 @@ func (cmd *WorkerCMD) Run(ctx *cliContext.Context) error {
 }
 
 // subscribeFileStaging subscribes to NATS file staging subjects for this node.
-func (cmd *WorkerCMD) subscribeFileStaging(natsClient *messaging.Client, nodeID string) error {
+func (cmd *WorkerCMD) subscribeFileStaging(natsClient messaging.MessagingClient, nodeID string) error {
 	// Create FileManager with same S3 config as the frontend
 	s3Store, err := storage.NewS3Store(storage.S3Config{
 		Endpoint:       cmd.StorageURL,
@@ -346,7 +349,7 @@ type backendSupervisor struct {
 	systemState *system.SystemState
 	galleries   []config.Gallery
 	nodeID      string
-	nats        *messaging.Client
+	nats        messaging.MessagingClient
 	sigCh       chan<- os.Signal // send shutdown signal instead of os.Exit
 
 	mu        sync.Mutex

core/cli/workerregistry/client.go

@@ -68,7 +68,7 @@ func (c *RegistrationClient) Register(ctx context.Context, body map[string]any)
 	}
 	defer resp.Body.Close()
 
-	if resp.StatusCode != http.StatusOK {
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
 		return "", "", fmt.Errorf("registration failed with status %d", resp.StatusCode)
 	}
 

core/config/distributed_config.go

@@ -2,7 +2,10 @@ package config
 
 import (
 	"cmp"
+	"fmt"
 	"time"
+
+	"github.com/mudler/xlog"
 )
 
 // DistributedConfig holds configuration for horizontal scaling mode.
@@ -28,13 +31,50 @@ 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)
+	PerModelHealthCheck bool          // Enable per-model backend health checking (default false)
+	MCPCIJobTimeout     time.Duration // MCP CI job execution timeout (default 10m)
 
 	MaxUploadSize int64 // Maximum upload body size in bytes (default 50 GB)
 
 	AgentWorkerConcurrency int `yaml:"agent_worker_concurrency" json:"agent_worker_concurrency" env:"LOCALAI_AGENT_WORKER_CONCURRENCY"`
 	JobWorkerConcurrency   int `yaml:"job_worker_concurrency" json:"job_worker_concurrency" env:"LOCALAI_JOB_WORKER_CONCURRENCY"`
 }
 
+// Validate checks that the distributed configuration is internally consistent.
+// It returns nil if distributed mode is disabled.
+func (c DistributedConfig) Validate() error {
+	if !c.Enabled {
+		return nil
+	}
+	if c.NatsURL == "" {
+		return fmt.Errorf("distributed mode requires --nats-url / LOCALAI_NATS_URL")
+	}
+	// S3 credentials must be paired
+	if (c.StorageAccessKey != "" && c.StorageSecretKey == "") ||
+		(c.StorageAccessKey == "" && c.StorageSecretKey != "") {
+		return fmt.Errorf("storage-access-key and storage-secret-key must both be set or both empty")
+	}
+	// Warn about missing registration token (not an error)
+	if c.RegistrationToken == "" {
+		xlog.Warn("distributed mode running without registration token — node endpoints are unprotected")
+	}
+	// Check for negative durations
+	for name, d := range map[string]time.Duration{
+		"mcp-tool-timeout":      c.MCPToolTimeout,
+		"mcp-discovery-timeout": c.MCPDiscoveryTimeout,
+		"worker-wait-timeout":   c.WorkerWaitTimeout,
+		"drain-timeout":         c.DrainTimeout,
+		"health-check-interval": c.HealthCheckInterval,
+		"stale-node-threshold":  c.StaleNodeThreshold,
+		"mcp-ci-job-timeout":    c.MCPCIJobTimeout,
+	} {
+		if d < 0 {
+			return fmt.Errorf("%s must not be negative", name)
+		}
+	}
+	return nil
+}
+
 // Distributed config options
 
 var EnableDistributed = func(o *ApplicationConfig) {
@@ -101,6 +141,7 @@ const (
 	DefaultDrainTimeout        = 30 * time.Second
 	DefaultHealthCheckInterval = 15 * time.Second
 	DefaultStaleNodeThreshold  = 60 * time.Second
+	DefaultMCPCIJobTimeout     = 10 * time.Minute
 )
 
 // DefaultMaxUploadSize is the default maximum upload body size (50 GB).
@@ -136,6 +177,11 @@ func (c DistributedConfig) StaleNodeThresholdOrDefault() time.Duration {
 	return cmp.Or(c.StaleNodeThreshold, DefaultStaleNodeThreshold)
 }
 
+// MCPCIJobTimeoutOrDefault returns the configured MCP CI job timeout or the default.
+func (c DistributedConfig) MCPCIJobTimeoutOrDefault() time.Duration {
+	return cmp.Or(c.MCPCIJobTimeout, DefaultMCPCIJobTimeout)
+}
+
 // MaxUploadSizeOrDefault returns the configured max upload size or the default.
 func (c DistributedConfig) MaxUploadSizeOrDefault() int64 {
 	if c.MaxUploadSize > 0 {

core/http/endpoints/anthropic/messages.go

@@ -61,15 +61,13 @@ func MessagesEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evalu
 			if mcpErr == nil {
 				mcpExecutor = mcpTools.NewToolExecutor(c.Request().Context(), natsClient, cfg.Name, remote, stdio, mcpServers)
 
-				// Prompt and resource injection (local mode only)
-				if natsClient == nil {
-					namedSessions, sessErr := mcpTools.NamedSessionsFromMCPConfig(cfg.Name, remote, stdio, mcpServers)
-					if sessErr == nil && len(namedSessions) > 0 {
-						mcpCtx, _ := mcpTools.InjectMCPContext(c.Request().Context(), namedSessions, mcpPromptName, mcpPromptArgs, mcpResourceURIs)
-						if mcpCtx != nil {
-							openAIMessages = append(mcpCtx.PromptMessages, openAIMessages...)
-							mcpTools.AppendResourceSuffix(openAIMessages, mcpCtx.ResourceSuffix)
-						}
+				// Prompt and resource injection (pre-processing step — resolves locally regardless of distributed mode)
+				namedSessions, sessErr := mcpTools.NamedSessionsFromMCPConfig(cfg.Name, remote, stdio, mcpServers)
+				if sessErr == nil && len(namedSessions) > 0 {
+					mcpCtx, _ := mcpTools.InjectMCPContext(c.Request().Context(), namedSessions, mcpPromptName, mcpPromptArgs, mcpResourceURIs)
+					if mcpCtx != nil {
+						openAIMessages = append(mcpCtx.PromptMessages, openAIMessages...)
+						mcpTools.AppendResourceSuffix(openAIMessages, mcpCtx.ResourceSuffix)
 					}
 				}