refactoring and consolidation

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

Author: mudler

core/application/distributed.go

@@ -13,6 +13,7 @@ import (
 	"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/pkg/sanitize"
 	"github.com/mudler/xlog"
 	"gorm.io/gorm"
 )
@@ -50,7 +51,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*Distribut
 		return nil, fmt.Errorf("distributed mode requires authentication to be enabled (--auth / LOCALAI_AUTH=true)")
 	}
 	if !isPostgresURL(cfg.Auth.DatabaseURL) {
-		return nil, fmt.Errorf("distributed mode requires PostgreSQL for auth database (got %q)", cfg.Auth.DatabaseURL)
+		return nil, fmt.Errorf("distributed mode requires PostgreSQL for auth database (got %q)", sanitize.URL(cfg.Auth.DatabaseURL))
 	}
 
 	// Validate NATS
@@ -69,7 +70,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*Distribut
 	if err != nil {
 		return nil, fmt.Errorf("connecting to NATS: %w", err)
 	}
-	xlog.Info("Connected to NATS", "url", cfg.Distributed.NatsURL)
+	xlog.Info("Connected to NATS", "url", sanitize.URL(cfg.Distributed.NatsURL))
 
 	// Ensure NATS is closed if any subsequent initialization step fails.
 	success := false
@@ -143,7 +144,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*Distribut
 	xlog.Info("Distributed job store initialized")
 
 	// Initialize job dispatcher
-	dispatcher := jobs.NewDispatcher(jobStore, natsClient, authDB, cfg.Distributed.InstanceID)
+	dispatcher := jobs.NewDispatcher(jobStore, natsClient, authDB, cfg.Distributed.InstanceID, cfg.Distributed.JobWorkerConcurrency)
 
 	// Initialize agent store
 	agentStore, err := agents.NewAgentStore(authDB)

core/application/startup.go

@@ -21,6 +21,7 @@ import (
 	"github.com/mudler/LocalAI/internal"
 
 	"github.com/mudler/LocalAI/pkg/model"
+	"github.com/mudler/LocalAI/pkg/sanitize"
 	"github.com/mudler/LocalAI/pkg/xsysinfo"
 	"github.com/mudler/xlog"
 )
@@ -104,7 +105,7 @@ func New(opts ...config.AppOption) (*Application, error) {
 			return nil, fmt.Errorf("failed to initialize auth database: %w", err)
 		}
 		application.authDB = authDB
-		xlog.Info("Auth enabled", "database", options.Auth.DatabaseURL)
+		xlog.Info("Auth enabled", "database", sanitize.URL(options.Auth.DatabaseURL))
 
 		// Start session and expired API key cleanup goroutine
 		go func() {

core/cli/agent_worker.go

@@ -19,6 +19,7 @@ import (
 	"github.com/mudler/LocalAI/core/services/jobs"
 	mcpRemote "github.com/mudler/LocalAI/core/services/mcp"
 	"github.com/mudler/LocalAI/core/services/messaging"
+	"github.com/mudler/LocalAI/pkg/sanitize"
 	"github.com/mudler/cogito"
 	"github.com/mudler/cogito/clients"
 	"github.com/mudler/xlog"
@@ -53,7 +54,7 @@ type AgentWorkerCMD struct {
 }
 
 func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
-	xlog.Info("Starting agent worker", "nats", cmd.NatsURL, "register_to", cmd.RegisterTo)
+	xlog.Info("Starting agent worker", "nats", sanitize.URL(cmd.NatsURL), "register_to", cmd.RegisterTo)
 
 	// Resolve API URL
 	apiURL := cmp.Or(cmd.APIURL, strings.TrimRight(cmd.RegisterTo, "/"))
@@ -118,11 +119,12 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
 	// Create and start the NATS dispatcher.
 	// No ConfigProvider or SkillStore needed — config and skills arrive in the job payload.
 	dispatcher := agents.NewNATSDispatcher(
-		&natsClientAdapter{natsClient},
+		natsClient,
 		eventBridge,
 		nil, // no ConfigProvider: config comes in the enriched NATS payload
 		apiURL, cmd.APIToken,
 		cmd.Subject, cmd.Queue,
+		0, // no concurrency limit (CLI worker)
 	)
 
 	if err := dispatcher.Start(shutdownCtx); err != nil {
@@ -273,21 +275,6 @@ func sendMCPDiscoveryReply(reply func([]byte), servers []mcpRemote.MCPServerInfo
 	reply(data)
 }
 
-// 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
-}
-
-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)) (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(shutdownCtx context.Context, data []byte, apiURL, apiToken string, natsClient *messaging.Client) {

core/cli/worker.go

@@ -24,6 +24,7 @@ import (
 	"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/pkg/sanitize"
 	grpc "github.com/mudler/LocalAI/pkg/grpc"
 	"github.com/mudler/LocalAI/pkg/model"
 	"github.com/mudler/LocalAI/pkg/system"
@@ -126,7 +127,7 @@ func (cmd *WorkerCMD) Run(ctx *cliContext.Context) error {
 	}
 
 	// Connect to NATS
-	xlog.Info("Connecting to NATS", "url", cmd.NatsURL)
+	xlog.Info("Connecting to NATS", "url", sanitize.URL(cmd.NatsURL))
 	natsClient, err := messaging.New(cmd.NatsURL)
 	if err != nil {
 		return fmt.Errorf("connecting to NATS: %w", err)

core/config/distributed_config.go

@@ -30,6 +30,9 @@ type DistributedConfig struct {
 	StaleNodeThreshold  time.Duration // Time before a node is considered stale (default 60s)
 
 	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"`
 }
 
 // Distributed config options

core/config/model_config.go

@@ -50,7 +50,7 @@ type ModelConfig struct {
 	InputToken                                 [][]int                `yaml:"-" json:"-"`
 	functionCallString, functionCallNameString string                 `yaml:"-" json:"-"`
 	ResponseFormat                             string                 `yaml:"-" json:"-"`
-	ResponseFormatMap                          map[string]interface{} `yaml:"-" json:"-"`
+	ResponseFormatMap                          map[string]any `yaml:"-" json:"-"`
 
 	FunctionsConfig functions.FunctionsConfig `yaml:"function,omitempty" json:"function,omitempty"`
 	ReasoningConfig reasoning.Config          `yaml:"reasoning,omitempty" json:"reasoning,omitempty"`

core/http/endpoints/localai/agents.go

@@ -4,10 +4,11 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"maps"
 	"net/http"
 	"os"
 	"path/filepath"
-	"sort"
+	"slices"
 	"strings"
 
 	"github.com/labstack/echo/v4"
@@ -76,11 +77,7 @@ func ListAgentsEndpoint(app *application.Application) echo.HandlerFunc {
 		svc := app.AgentPoolService()
 		userID := getUserID(c)
 		statuses := svc.ListAgentsForUser(userID)
-		agents := make([]string, 0, len(statuses))
-		for name := range statuses {
-			agents = append(agents, name)
-		}
-		sort.Strings(agents)
+		agents := slices.Sorted(maps.Keys(statuses))
 		resp := map[string]any{
 			"agents":     agents,
 			"agentCount": len(agents),

core/http/endpoints/localai/nodes.go

@@ -211,7 +211,9 @@ func provisionAgentWorkerKey(authDB *gorm.DB, registry *nodes.NodeRegistry, node
 
 	node.AuthUserID = workerUser.ID
 	node.APIKeyID = apiKey.ID
-	registry.UpdateAuthRefs(node.ID, workerUser.ID, apiKey.ID)
+	if err := registry.UpdateAuthRefs(node.ID, workerUser.ID, apiKey.ID); err != nil {
+		xlog.Warn("Failed to update auth refs on node", "node", node.Name, "error", err)
+	}
 
 	// Grant collections feature so the worker can store/retrieve KB data on behalf of users.
 	perm := &auth.UserPermission{

core/http/routes/nodes.go

@@ -29,10 +29,10 @@ func nodeReadyMiddleware(registry *nodes.NodeRegistry) echo.MiddlewareFunc {
 // 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.
+// TODO(security): Node self-service endpoints authenticate via shared registration
+// token but do not verify per-node identity. A compromised worker can heartbeat/drain/
+// deregister other nodes. Future: issue per-node JWT at registration, validate node
+// identity on subsequent requests (compare :id param with token subject).
 func RegisterNodeSelfServiceRoutes(e *echo.Echo, registry *nodes.NodeRegistry, registrationToken string, autoApprove bool, authDB *gorm.DB, hmacSecret string) {
 	if registry == nil {
 		return
@@ -84,9 +84,14 @@ func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloade
 	e.GET("/ws/nodes/:id/backend-logs/:modelId", localai.NodeBackendLogsWSEndpoint(registry, registrationToken), readyMw, adminMw)
 }
 
-// nodeTokenAuth returns middleware that validates the registration token
-// from an Authorization: Bearer <token> header using constant-time comparison.
-// If registrationToken is empty, all requests are allowed (no auth required).
+// nodeTokenAuth validates the registration token for node self-service endpoints.
+// When registrationToken is empty (single-node / non-distributed mode), these
+// endpoints are unprotected. This is intentional: in single-node mode there are
+// no remote workers to authenticate. Operators enabling distributed mode MUST
+// set a registration token via LOCALAI_REGISTRATION_TOKEN or config.
+//
+// It validates the token from an Authorization: Bearer <token> header using
+// constant-time comparison.
 func nodeTokenAuth(registrationToken string) echo.MiddlewareFunc {
 	return func(next echo.HandlerFunc) echo.HandlerFunc {
 		return func(c echo.Context) error {

core/http/routes/ui_api.go

@@ -9,7 +9,8 @@ import (
 	"net/http"
 	"net/url"
 	"path"
-	"sort"
+	"cmp"
+	"slices"
 	"strconv"
 	"strings"
 	"sync"
@@ -153,17 +154,17 @@ func RegisterUIAPIRoutes(app *echo.Echo, cl *config.ModelConfigLoader, ml *model
 		}
 
 		// Sort operations by progress (ascending), then by ID for stable display order
-		sort.Slice(operations, func(i, j int) bool {
-			progressI := operations[i]["progress"].(int)
-			progressJ := operations[j]["progress"].(int)
+		slices.SortFunc(operations, func(a, b map[string]interface{}) int {
+			progressA := a["progress"].(int)
+			progressB := b["progress"].(int)
 
 			// Primary sort by progress
-			if progressI != progressJ {
-				return progressI < progressJ
+			if progressA != progressB {
+				return cmp.Compare(progressA, progressB)
 			}
 
 			// Secondary sort by ID for stability when progress is the same
-			return operations[i]["id"].(string) < operations[j]["id"].(string)
+			return cmp.Compare(a["id"].(string), b["id"].(string))
 		})
 
 		return c.JSON(200, map[string]interface{}{
@@ -239,7 +240,7 @@ func RegisterUIAPIRoutes(app *echo.Echo, cl *config.ModelConfigLoader, ml *model
 		for t := range allTags {
 			tags = append(tags, t)
 		}
-		sort.Strings(tags)
+		slices.Sort(tags)
 
 		// Get all available backends (before filtering so dropdown always shows all)
 		allBackendsMap := map[string]struct{}{}
@@ -252,7 +253,7 @@ func RegisterUIAPIRoutes(app *echo.Echo, cl *config.ModelConfigLoader, ml *model
 		for b := range allBackendsMap {
 			backendNames = append(backendNames, b)
 		}
-		sort.Strings(backendNames)
+		slices.Sort(backendNames)
 
 		if tag != "" {
 			models = gallery.GalleryElements[*gallery.GalleryModel](models).FilterByTag(tag)
@@ -809,7 +810,7 @@ func RegisterUIAPIRoutes(app *echo.Echo, cl *config.ModelConfigLoader, ml *model
 		for t := range allTags {
 			tags = append(tags, t)
 		}
-		sort.Strings(tags)
+		slices.Sort(tags)
 
 		if tag != "" {
 			backends = gallery.GalleryElements[*gallery.GalleryBackend](backends).FilterByTag(tag)