enhancements, refactorings

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

Author: mudler

core/application/application.go

@@ -167,26 +167,27 @@ func (a *Application) StartAgentPool() {
 	if !a.applicationConfig.AgentPool.Enabled {
 		return
 	}
-	aps, err := agentpool.NewAgentPoolService(a.applicationConfig)
+	// Build options struct from available dependencies
+	opts := agentpool.AgentPoolOptions{
+		AuthDB: a.authDB,
+	}
+	if d := a.Distributed(); d != nil {
+		if d.DistStores != nil && d.DistStores.Skills != nil {
+			opts.SkillStore = d.DistStores.Skills
+		}
+		opts.NATSClient = d.Nats
+		opts.EventBridge = d.AgentBridge
+		opts.AgentStore = d.AgentStore
+	}
+
+	aps, err := agentpool.NewAgentPoolService(a.applicationConfig, opts)
 	if err != nil {
 		xlog.Error("Failed to create agent pool service", "error", err)
 		return
 	}
-	if a.authDB != nil {
-		aps.SetAuthDB(a.authDB)
-	}
+
 	// Wire distributed mode components
 	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 {

core/application/distributed.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"strings"
 
 	"github.com/google/uuid"
@@ -49,6 +50,9 @@ func (ds *DistributedServices) Shutdown() {
 	if ds.Dispatcher != nil {
 		ds.Dispatcher.Stop()
 	}
+	if closer, ok := ds.Store.(io.Closer); ok {
+		closer.Close()
+	}
 	// AgentBridge has no Close method — its NATS subscriptions are cleaned up
 	// when the NATS client is closed below.
 	if ds.Nats != nil {
@@ -107,7 +111,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB) (*Distribut
 		if cfg.Distributed.StorageBucket == "" {
 			return nil, fmt.Errorf("distributed storage bucket must be set when storage URL is configured")
 		}
-		s3Store, err := storage.NewS3Store(storage.S3Config{
+		s3Store, err := storage.NewS3Store(context.Background(), storage.S3Config{
 			Endpoint:        cfg.Distributed.StorageURL,
 			Region:          cfg.Distributed.StorageRegion,
 			Bucket:          cfg.Distributed.StorageBucket,

core/cli/worker.go

@@ -189,7 +189,8 @@ func (cmd *WorkerCMD) Run(ctx *cliContext.Context) error {
 // subscribeFileStaging subscribes to NATS file staging subjects for this node.
 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{
+	// TODO: propagate a caller-provided context once WorkerCMD carries one
+	s3Store, err := storage.NewS3Store(context.Background(), storage.S3Config{
 		Endpoint:       cmd.StorageURL,
 		Region:         cmd.StorageRegion,
 		Bucket:         cmd.StorageBucket,
@@ -383,9 +384,10 @@ func (s *backendSupervisor) startBackend(backend, backendPath string) (string, e
 		port = s.nextPort
 		s.nextPort++
 	}
-	addr := fmt.Sprintf("0.0.0.0:%d", port)
+	bindAddr := fmt.Sprintf("0.0.0.0:%d", port)
+	clientAddr := fmt.Sprintf("127.0.0.1:%d", port)
 
-	proc, err := s.ml.StartProcess(backendPath, backend, addr)
+	proc, err := s.ml.StartProcess(backendPath, backend, bindAddr)
 	if err != nil {
 		s.mu.Unlock()
 		return "", fmt.Errorf("starting backend process: %w", err)
@@ -394,17 +396,17 @@ func (s *backendSupervisor) startBackend(backend, backendPath string) (string, e
 	s.processes[backend] = &backendProcess{
 		proc:    proc,
 		backend: backend,
-		addr:    addr,
+		addr:    clientAddr,
 	}
-	xlog.Info("Backend process started", "backend", backend, "addr", addr)
+	xlog.Info("Backend process started", "backend", backend, "addr", clientAddr)
 
 	// Capture reference before unlocking for race-safe health check.
 	// Another goroutine could stopBackend and recycle the port while we poll.
 	bp := s.processes[backend]
 	s.mu.Unlock()
 
 	// Wait for the gRPC server to be ready
-	client := grpc.NewClientWithToken(addr, false, nil, false, s.cmd.RegistrationToken)
+	client := grpc.NewClientWithToken(clientAddr, false, nil, false, s.cmd.RegistrationToken)
 	for range 20 {
 		time.Sleep(200 * time.Millisecond)
 		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
@@ -417,27 +419,34 @@ func (s *backendSupervisor) startBackend(backend, backendPath string) (string, e
 			if !exists || currentBP != bp {
 				return "", fmt.Errorf("backend %s was stopped during startup", backend)
 			}
-			xlog.Debug("Backend gRPC server is ready", "backend", backend, "addr", addr)
-			return addr, nil
+			xlog.Debug("Backend gRPC server is ready", "backend", backend, "addr", clientAddr)
+			return clientAddr, nil
 		}
 		cancel()
 	}
 
-	xlog.Warn("Backend gRPC server not ready after waiting, proceeding anyway", "backend", backend, "addr", addr)
-	return addr, nil
+	xlog.Warn("Backend gRPC server not ready after waiting, proceeding anyway", "backend", backend, "addr", clientAddr)
+	return clientAddr, nil
 }
 
 // stopBackend stops a specific backend's gRPC process.
 func (s *backendSupervisor) stopBackend(backend string) {
 	s.mu.Lock()
-	defer s.mu.Unlock()
-
 	bp, ok := s.processes[backend]
 	if !ok || bp.proc == nil {
+		s.mu.Unlock()
 		return
 	}
+	// Clean up map and recycle port while holding lock
+	delete(s.processes, backend)
+	if _, portStr, err := net.SplitHostPort(bp.addr); err == nil {
+		if p, err := strconv.Atoi(portStr); err == nil {
+			s.freePorts = append(s.freePorts, p)
+		}
+	}
+	s.mu.Unlock()
 
-	// Best-effort Free() to release GPU memory
+	// Network I/O outside the lock
 	client := grpc.NewClientWithToken(bp.addr, false, nil, false, s.cmd.RegistrationToken)
 	if freeFunc, ok := client.(interface{ Free() error }); ok {
 		xlog.Debug("Calling Free() before stopping backend", "backend", backend)
@@ -450,13 +459,6 @@ func (s *backendSupervisor) stopBackend(backend string) {
 	if err := bp.proc.Stop(); err != nil {
 		xlog.Error("Error stopping backend process", "backend", backend, "error", err)
 	}
-	// Recycle the port for future backends
-	if _, portStr, err := net.SplitHostPort(bp.addr); err == nil {
-		if p, err := strconv.Atoi(portStr); err == nil {
-			s.freePorts = append(s.freePorts, p)
-		}
-	}
-	delete(s.processes, backend)
 }
 
 // stopAllBackends stops all running backend processes.

core/cli/workerregistry/client.go

@@ -21,6 +21,7 @@ type RegistrationClient struct {
 	FrontendURL       string
 	RegistrationToken string
 	HTTPTimeout       time.Duration // used for registration calls; defaults to 10s
+	client            *http.Client
 }
 
 // httpTimeout returns the configured timeout or a sensible default.
@@ -31,6 +32,14 @@ func (c *RegistrationClient) httpTimeout() time.Duration {
 	return 10 * time.Second
 }
 
+// httpClient returns the shared HTTP client, initializing it on first use.
+func (c *RegistrationClient) httpClient() *http.Client {
+	if c.client == nil {
+		c.client = &http.Client{Timeout: c.httpTimeout()}
+	}
+	return c.client
+}
+
 // baseURL returns FrontendURL with any trailing slash stripped.
 func (c *RegistrationClient) baseURL() string {
 	return strings.TrimRight(c.FrontendURL, "/")
@@ -62,7 +71,7 @@ func (c *RegistrationClient) Register(ctx context.Context, body map[string]any)
 	req.Header.Set("Content-Type", "application/json")
 	c.setAuth(req)
 
-	resp, err := (&http.Client{Timeout: c.httpTimeout()}).Do(req)
+	resp, err := c.httpClient().Do(req)
 	if err != nil {
 		return "", "", fmt.Errorf("posting to %s: %w", url, err)
 	}
@@ -114,11 +123,11 @@ func (c *RegistrationClient) Heartbeat(ctx context.Context, nodeID string, body
 	req.Header.Set("Content-Type", "application/json")
 	c.setAuth(req)
 
-	resp, err := (&http.Client{Timeout: 5 * time.Second}).Do(req)
+	resp, err := c.httpClient().Do(req)
 	if err != nil {
 		return err
 	}
-	resp.Body.Close()
+	defer resp.Body.Close()
 	return nil
 }
 
@@ -147,11 +156,11 @@ func (c *RegistrationClient) Drain(ctx context.Context, nodeID string) error {
 	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, url, nil)
 	c.setAuth(req)
 
-	resp, err := (&http.Client{Timeout: 5 * time.Second}).Do(req)
+	resp, err := c.httpClient().Do(req)
 	if err != nil {
 		return err
 	}
-	resp.Body.Close()
+	defer resp.Body.Close()
 	if resp.StatusCode != http.StatusOK {
 		return fmt.Errorf("drain failed with status %d", resp.StatusCode)
 	}
@@ -162,16 +171,17 @@ func (c *RegistrationClient) Drain(ctx context.Context, nodeID string) error {
 // in-flight requests, or until timeout elapses.
 func (c *RegistrationClient) WaitForDrain(ctx context.Context, nodeID string, timeout time.Duration) {
 	url := c.baseURL() + "/api/node/" + nodeID + "/models"
-	client := &http.Client{Timeout: 5 * time.Second}
 
 	deadline := time.Now().Add(timeout)
 	for time.Now().Before(deadline) {
 		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
 		c.setAuth(req)
 
-		resp, err := client.Do(req)
+		resp, err := c.httpClient().Do(req)
 		if err != nil {
-			break
+			xlog.Warn("Drain poll failed, will retry", "error", err)
+			time.Sleep(1 * time.Second)
+			continue
 		}
 		var models []struct {
 			InFlight int `json:"in_flight"`
@@ -201,11 +211,11 @@ func (c *RegistrationClient) Deregister(ctx context.Context, nodeID string) erro
 	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, url, nil)
 	c.setAuth(req)
 
-	resp, err := (&http.Client{Timeout: 5 * time.Second}).Do(req)
+	resp, err := c.httpClient().Do(req)
 	if err != nil {
 		return err
 	}
-	resp.Body.Close()
+	defer resp.Body.Close()
 	if resp.StatusCode != http.StatusOK {
 		return fmt.Errorf("deregistration failed with status %d", resp.StatusCode)
 	}

core/services/advisorylock/advisorylock_test.go

@@ -10,7 +10,6 @@ import (
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 
-	"github.com/mudler/LocalAI/core/services/messaging"
 	"github.com/mudler/LocalAI/core/services/testutil"
 	"gorm.io/gorm"
 )
@@ -97,12 +96,12 @@ var _ = Describe("AdvisoryLock", func() {
 
 		It("all advisory lock keys are unique", func() {
 			keys := map[int64]string{
-				messaging.AdvisoryLockCronScheduler:    "messaging.AdvisoryLockCronScheduler",
-				messaging.AdvisoryLockStaleNodeCleanup: "messaging.AdvisoryLockStaleNodeCleanup",
-				messaging.AdvisoryLockGalleryDedup:     "messaging.AdvisoryLockGalleryDedup",
-				messaging.AdvisoryLockAgentScheduler:   "messaging.AdvisoryLockAgentScheduler",
-				messaging.AdvisoryLockHealthCheck:      "messaging.AdvisoryLockHealthCheck",
-				messaging.AdvisoryLockSchemaMigrate:    "messaging.AdvisoryLockSchemaMigrate",
+				KeyCronScheduler:    "KeyCronScheduler",
+				KeyStaleNodeCleanup: "KeyStaleNodeCleanup",
+				KeyGalleryDedup:     "KeyGalleryDedup",
+				KeyAgentScheduler:   "KeyAgentScheduler",
+				KeyHealthCheck:      "KeyHealthCheck",
+				KeySchemaMigrate:    "KeySchemaMigrate",
 			}
 
 			Expect(keys).To(HaveLen(6), "some advisory lock keys have the same value")

core/services/advisorylock/keys.go

@@ -0,0 +1,13 @@
+package advisorylock
+
+// Advisory lock keys for distributed coordination via PostgreSQL.
+// These keys are global across the database — avoid collisions with
+// other applications sharing the same PostgreSQL instance.
+const (
+	KeyCronScheduler    int64 = 100
+	KeyStaleNodeCleanup int64 = 101
+	KeyGalleryDedup     int64 = 102
+	KeyAgentScheduler   int64 = 103
+	KeyHealthCheck      int64 = 104
+	KeySchemaMigrate    int64 = 105
+)

core/services/agentpool/agent_pool.go

@@ -79,8 +79,8 @@ type AgentEventBridge interface {
 	PublishMessage(agentName, userID, sender, content, messageID string) error
 	PublishStatus(agentName, userID, status string) error
 	PublishStreamEvent(agentName, userID string, data map[string]any) error
-	RegisterCancel(agentName, userID string, cancel context.CancelFunc)
-	DeregisterCancel(agentName, userID string)
+	RegisterCancel(key string, cancel context.CancelFunc)
+	DeregisterCancel(key string)
 }
 
 // AgentConfigStore is the interface for agent config persistence.
@@ -93,10 +93,39 @@ type AgentConfigStore interface {
 	UpdateLastRun(userID, name string) error
 }
 
-func NewAgentPoolService(appConfig *config.ApplicationConfig) (*AgentPoolService, error) {
-	return &AgentPoolService{
+// AgentPoolOptions holds optional dependencies for AgentPoolService.
+// Zero values are fine — the service degrades gracefully without them.
+type AgentPoolOptions struct {
+	AuthDB      *gorm.DB
+	SkillStore  *distributed.SkillStore
+	NATSClient  messaging.Publisher
+	EventBridge AgentEventBridge
+	AgentStore  *agents.AgentStore
+}
+
+func NewAgentPoolService(appConfig *config.ApplicationConfig, opts ...AgentPoolOptions) (*AgentPoolService, error) {
+	svc := &AgentPoolService{
 		appConfig: appConfig,
-	}, nil
+	}
+	if len(opts) > 0 {
+		o := opts[0]
+		if o.AuthDB != nil {
+			svc.users.authDB = o.AuthDB
+		}
+		if o.SkillStore != nil {
+			svc.distributed.skillStore = o.SkillStore
+		}
+		if o.NATSClient != nil {
+			svc.distributed.natsClient = o.NATSClient
+		}
+		if o.EventBridge != nil {
+			svc.distributed.eventBridge = o.EventBridge
+		}
+		if o.AgentStore != nil {
+			svc.distributed.agentStore = o.AgentStore
+		}
+	}
+	return svc, nil
 }
 
 func (s *AgentPoolService) Start(ctx context.Context) error {
@@ -337,16 +366,19 @@ func (s *AgentPoolService) Pool() *state.AgentPool {
 }
 
 // SetNATSClient sets the NATS client for distributed agent execution.
+// Deprecated: prefer passing NATSClient via AgentPoolOptions at construction time.
 func (s *AgentPoolService) SetNATSClient(nc messaging.Publisher) {
 	s.distributed.natsClient = nc
 }
 
 // SetEventBridge sets the event bridge for distributed SSE + persistence.
+// Deprecated: prefer passing EventBridge via AgentPoolOptions at construction time.
 func (s *AgentPoolService) SetEventBridge(eb AgentEventBridge) {
 	s.distributed.eventBridge = eb
 }
 
 // SetAgentStore sets the PostgreSQL agent config store.
+// Deprecated: prefer passing AgentStore via AgentPoolOptions at construction time.
 func (s *AgentPoolService) SetAgentStore(store *agents.AgentStore) {
 	s.distributed.agentStore = store
 }
@@ -578,11 +610,13 @@ func (s *AgentPoolService) UserServicesManager() *UserServicesManager {
 }
 
 // SetAuthDB sets the auth database for API key generation.
+// Deprecated: prefer passing AuthDB via AgentPoolOptions at construction time.
 func (s *AgentPoolService) SetAuthDB(db *gorm.DB) {
 	s.users.authDB = db
 }
 
 // SetSkillStore sets the distributed skill store for persisting skill metadata to PostgreSQL.
+// Deprecated: prefer passing SkillStore via AgentPoolOptions at construction time.
 func (s *AgentPoolService) SetSkillStore(store *distributed.SkillStore) {
 	s.distributed.skillStore = store
 }