refactoring and consolidation

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

Author: mudler

core/cli/agent_worker.go

@@ -139,14 +139,18 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
 
 	// Subscribe to MCP tool execution requests (load-balanced across workers).
 	// The frontend routes model-level MCP tool calls here via NATS request-reply.
-	natsClient.QueueSubscribeReply(messaging.SubjectMCPToolExecute, messaging.QueueAgentWorkers, func(data []byte, reply func([]byte)) {
+	if _, err := natsClient.QueueSubscribeReply(messaging.SubjectMCPToolExecute, messaging.QueueAgentWorkers, func(data []byte, reply func([]byte)) {
 		handleMCPToolRequest(data, reply)
-	})
+	}); err != nil {
+		return fmt.Errorf("subscribing to %s: %w", messaging.SubjectMCPToolExecute, err)
+	}
 
 	// Subscribe to MCP discovery requests (load-balanced across workers).
-	natsClient.QueueSubscribeReply(messaging.SubjectMCPDiscovery, messaging.QueueAgentWorkers, func(data []byte, reply func([]byte)) {
+	if _, err := natsClient.QueueSubscribeReply(messaging.SubjectMCPDiscovery, messaging.QueueAgentWorkers, func(data []byte, reply func([]byte)) {
 		handleMCPDiscoveryRequest(data, reply)
-	})
+	}); err != nil {
+		return fmt.Errorf("subscribing to %s: %w", messaging.SubjectMCPDiscovery, err)
+	}
 
 	// Subscribe to MCP CI job execution (load-balanced across agent workers).
 	// In distributed mode, MCP CI jobs are routed here because the frontend
@@ -157,21 +161,25 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
 	}
 	mcpCIJobTimeout = cmp.Or(mcpCIJobTimeout, config.DefaultMCPCIJobTimeout)
 
-	natsClient.QueueSubscribe(messaging.SubjectMCPCIJobsNew, messaging.QueueWorkers, func(data []byte) {
+	if _, err := natsClient.QueueSubscribe(messaging.SubjectMCPCIJobsNew, messaging.QueueWorkers, func(data []byte) {
 		handleMCPCIJob(shutdownCtx, data, apiURL, cmd.APIToken, natsClient, mcpCIJobTimeout)
-	})
+	}); err != nil {
+		return fmt.Errorf("subscribing to %s: %w", messaging.SubjectMCPCIJobsNew, err)
+	}
 
 	// Subscribe to backend stop events to clean up cached MCP sessions.
 	// In the main application this is done via ml.OnModelUnload, but the agent
 	// worker has no model loader — we listen for the NATS stop event instead.
-	natsClient.Subscribe(messaging.SubjectNodeBackendStop(nodeID), func(data []byte) {
+	if _, err := natsClient.Subscribe(messaging.SubjectNodeBackendStop(nodeID), func(data []byte) {
 		var req struct {
 			Backend string `json:"backend"`
 		}
 		if json.Unmarshal(data, &req) == nil && req.Backend != "" {
 			mcpTools.CloseMCPSessions(req.Backend)
 		}
-	})
+	}); err != nil {
+		return fmt.Errorf("subscribing to %s: %w", messaging.SubjectNodeBackendStop(nodeID), err)
+	}
 
 	xlog.Info("Agent worker ready, waiting for jobs", "subject", cmd.Subject, "queue", cmd.Queue)
 

core/config/distributed_config.go

@@ -184,8 +184,5 @@ func (c DistributedConfig) MCPCIJobTimeoutOrDefault() time.Duration {
 
 // MaxUploadSizeOrDefault returns the configured max upload size or the default.
 func (c DistributedConfig) MaxUploadSizeOrDefault() int64 {
-	if c.MaxUploadSize > 0 {
-		return c.MaxUploadSize
-	}
-	return DefaultMaxUploadSize
+	return cmp.Or(c.MaxUploadSize, DefaultMaxUploadSize)
 }

core/http/auth/middleware.go

@@ -545,6 +545,14 @@ func isAPIPath(path string) bool {
 		strings.HasPrefix(path, "/system") ||
 		strings.HasPrefix(path, "/ws/") ||
 		strings.HasPrefix(path, "/generated-") ||
+		strings.HasPrefix(path, "/chat/") ||
+		strings.HasPrefix(path, "/completions") ||
+		strings.HasPrefix(path, "/edits") ||
+		strings.HasPrefix(path, "/embeddings") ||
+		strings.HasPrefix(path, "/audio/") ||
+		strings.HasPrefix(path, "/images/") ||
+		strings.HasPrefix(path, "/messages") ||
+		strings.HasPrefix(path, "/responses") ||
 		path == "/metrics"
 }
 

core/http/endpoints/localai/cors_proxy.go

@@ -3,6 +3,7 @@ package localai
 import (
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/url"
 	"strings"
@@ -13,6 +14,33 @@ import (
 	"github.com/mudler/xlog"
 )
 
+var privateNetworks []*net.IPNet
+
+func init() {
+	for _, cidr := range []string{
+		"10.0.0.0/8",
+		"172.16.0.0/12",
+		"192.168.0.0/16",
+		"127.0.0.0/8",
+		"169.254.0.0/16",
+		"::1/128",
+		"fc00::/7",
+		"fe80::/10",
+	} {
+		_, network, _ := net.ParseCIDR(cidr)
+		privateNetworks = append(privateNetworks, network)
+	}
+}
+
+func isPrivateIP(ip net.IP) bool {
+	for _, network := range privateNetworks {
+		if network.Contains(ip) {
+			return true
+		}
+	}
+	return false
+}
+
 var corsProxyClient = &http.Client{
 	Timeout: 10 * time.Minute,
 }
@@ -36,6 +64,16 @@ func CORSProxyEndpoint(appConfig *config.ApplicationConfig) echo.HandlerFunc {
 			return c.JSON(http.StatusBadRequest, map[string]string{"error": "only http and https schemes are supported"})
 		}
 
+		ips, err := net.LookupIP(parsed.Hostname())
+		if err != nil {
+			return c.JSON(http.StatusBadRequest, map[string]string{"error": "cannot resolve hostname"})
+		}
+		for _, ip := range ips {
+			if isPrivateIP(ip) {
+				return c.JSON(http.StatusForbidden, map[string]string{"error": "requests to private networks are not allowed"})
+			}
+		}
+
 		xlog.Debug("CORS proxy request", "method", c.Request().Method, "target", targetURL)
 
 		proxyReq, err := http.NewRequestWithContext(
@@ -53,6 +91,8 @@ func CORSProxyEndpoint(appConfig *config.ApplicationConfig) echo.HandlerFunc {
 			"Host": true, "Connection": true, "Keep-Alive": true,
 			"Transfer-Encoding": true, "Upgrade": true, "Origin": true,
 			"Referer": true,
+			"Authorization": true, "Cookie": true,
+			"X-Api-Key": true, "Proxy-Authorization": true,
 		}
 		for key, values := range c.Request().Header {
 			if skipHeaders[key] {

core/http/routes/anthropic.go

@@ -75,14 +75,11 @@ func setAnthropicRequestContext(appConfig *config.ApplicationConfig) echo.Middle
 			reqCtx := c.Request().Context()
 			c1, cancel := context.WithCancel(appConfig.Context)
 
-			// Cancel when request context is cancelled (client disconnects)
-			go func() {
-				select {
-				case <-reqCtx.Done():
-					cancel()
-				case <-c1.Done():
-					// Already cancelled
-				}
+			// Bridge request cancellation to c1 without spawning a goroutine.
+			stop := context.AfterFunc(reqCtx, cancel)
+			defer func() {
+				stop()   // deregister callback if it hasn't fired
+				cancel() // release c1 resources (idempotent)
 			}()
 
 			// Add the correlation ID to the new context

core/http/routes/ui.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"net/url"
 	"slices"
 	"sync"
 	"time"
@@ -21,7 +22,15 @@ import (
 
 var backendLogsUpgrader = websocket.Upgrader{
 	CheckOrigin: func(r *http.Request) bool {
-		return true
+		origin := r.Header.Get("Origin")
+		if origin == "" {
+			return true // no origin header = same-origin or non-browser
+		}
+		u, err := url.Parse(origin)
+		if err != nil {
+			return false
+		}
+		return u.Host == r.Host
 	},
 }
 

core/services/advisorylock/leader_loop.go

@@ -21,7 +21,7 @@ func RunLeaderLoop(ctx context.Context, db *gorm.DB, lockKey int64, interval tim
 		case <-ctx.Done():
 			return
 		case <-ticker.C:
-			_, err := TryWithLock(db, lockKey, func() error {
+			_, err := TryWithLockCtx(ctx, db, lockKey, func() error {
 				fn()
 				return nil
 			})