enhancements, refactorings

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

Author: mudler

core/cli/worker.go

@@ -33,6 +33,29 @@ import (
 	"github.com/mudler/xlog"
 )
 
+// isPathAllowed checks if path is within one of the allowed directories.
+func isPathAllowed(path string, allowedDirs []string) bool {
+	absPath, err := filepath.Abs(path)
+	if err != nil {
+		return false
+	}
+	resolved, err := filepath.EvalSymlinks(absPath)
+	if err != nil {
+		// Path may not exist yet; use the absolute path
+		resolved = absPath
+	}
+	for _, dir := range allowedDirs {
+		absDir, err := filepath.Abs(dir)
+		if err != nil {
+			continue
+		}
+		if strings.HasPrefix(resolved, absDir+string(filepath.Separator)) || resolved == absDir {
+			return true
+		}
+	}
+	return false
+}
+
 // WorkerCMD starts a generic worker process for distributed mode.
 // Workers are backend-agnostic — they wait for backend.install NATS events
 // from the SmartRouter to install and start the required backend.
@@ -118,8 +141,6 @@ func (cmd *WorkerCMD) Run(ctx *cliContext.Context) error {
 	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
 	defer shutdownCancel()
 
-	go regClient.HeartbeatLoop(shutdownCtx, nodeID, heartbeatInterval, cmd.heartbeatBody)
-
 	// Start HTTP file transfer server
 	httpAddr := cmd.resolveHTTPAddr()
 	stagingDir := filepath.Join(cmd.ModelsPath, "..", "staging")
@@ -138,6 +159,27 @@ func (cmd *WorkerCMD) Run(ctx *cliContext.Context) error {
 	}
 	defer natsClient.Close()
 
+	// Start heartbeat goroutine (after NATS is connected so IsConnected check works)
+	go func() {
+		ticker := time.NewTicker(heartbeatInterval)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-shutdownCtx.Done():
+				return
+			case <-ticker.C:
+				if !natsClient.IsConnected() {
+					xlog.Warn("Skipping heartbeat: NATS disconnected")
+					continue
+				}
+				body := cmd.heartbeatBody()
+				if err := regClient.Heartbeat(shutdownCtx, nodeID, body); err != nil {
+					xlog.Warn("Heartbeat failed", "error", err)
+				}
+			}
+		}
+	}()
+
 	// Process supervisor — manages multiple backend gRPC processes on different ports
 	basePort := 50051
 	if cmd.Addr != "" {
@@ -242,6 +284,15 @@ func (cmd *WorkerCMD) subscribeFileStaging(natsClient messaging.MessagingClient,
 			return
 		}
 
+		allowedDirs := []string{cacheDir}
+		if cmd.ModelsPath != "" {
+			allowedDirs = append(allowedDirs, cmd.ModelsPath)
+		}
+		if !isPathAllowed(req.LocalPath, allowedDirs) {
+			replyJSON(reply, map[string]string{"error": "path outside allowed directories"})
+			return
+		}
+
 		if err := fm.Upload(context.Background(), req.Key, req.LocalPath); err != nil {
 			xlog.Error("File stage failed", "path", req.LocalPath, "key", req.Key, "error", err)
 			replyJSON(reply, map[string]string{"error": err.Error()})
@@ -450,9 +501,9 @@ func (s *backendSupervisor) stopBackend(backend string) {
 
 	// 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 {
+	if freeFunc, ok := client.(interface{ Free(context.Context) error }); ok {
 		xlog.Debug("Calling Free() before stopping backend", "backend", backend)
-		if err := freeFunc.Free(); err != nil {
+		if err := freeFunc.Free(context.Background()); err != nil {
 			xlog.Warn("Free() failed (best-effort)", "backend", backend, "error", err)
 		}
 	}
@@ -698,8 +749,8 @@ func (s *backendSupervisor) subscribeLifecycleEvents() {
 		if targetAddr != "" {
 			// Best-effort gRPC Free()
 			client := grpc.NewClientWithToken(targetAddr, false, nil, false, s.cmd.RegistrationToken)
-			if freeFunc, ok := client.(interface{ Free() error }); ok {
-				if err := freeFunc.Free(); err != nil {
+			if freeFunc, ok := client.(interface{ Free(context.Context) error }); ok {
+				if err := freeFunc.Free(context.Background()); err != nil {
 					xlog.Warn("Free() failed during model.unload", "error", err, "addr", targetAddr)
 				}
 			}

core/http/endpoints/anthropic/messages.go

@@ -434,7 +434,14 @@ func handleAnthropicStream(c echo.Context, id string, input *schema.AnthropicReq
 		_, tokenUsage, chatDeltas, err := openaiEndpoint.ComputeChoices(openAIReq, predInput, cfg, cl, appConfig, ml, func(s string, c *[]schema.Choice) {}, tokenCallback)
 		if err != nil {
 			xlog.Error("Anthropic stream model inference failed", "error", err)
-			return sendAnthropicError(c, 500, "api_error", fmt.Sprintf("model inference failed: %v", err))
+			sendAnthropicSSE(c, schema.AnthropicStreamEvent{
+				Type: "error",
+				Error: &schema.AnthropicError{
+					Type:    "api_error",
+					Message: fmt.Sprintf("model inference failed: %v", err),
+				},
+			})
+			return nil
 		}
 
 		// Also check chat deltas for tool calls

core/http/endpoints/openai/chat.go

@@ -59,10 +59,7 @@ func mergeToolCallDeltas(existing []schema.ToolCall, deltas []schema.ToolCall) [
 // @Success 200 {object} schema.OpenAIResponse "Response"
 // @Router /v1/chat/completions [post]
 func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, startupOptions *config.ApplicationConfig, natsClient mcpTools.MCPNATSClient) echo.HandlerFunc {
-	var id, textContentToReturn string
-	var created int
-
-	process := func(s string, req *schema.OpenAIRequest, config *config.ModelConfig, loader *model.ModelLoader, responses chan schema.OpenAIResponse, extraUsage bool) error {
+	process := func(s string, req *schema.OpenAIRequest, config *config.ModelConfig, loader *model.ModelLoader, responses chan schema.OpenAIResponse, extraUsage bool, id string, created int) error {
 		initialMessage := schema.OpenAIResponse{
 			ID:      id,
 			Created: created,
@@ -119,7 +116,7 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
 		close(responses)
 		return err
 	}
-	processTools := func(noAction string, prompt string, req *schema.OpenAIRequest, config *config.ModelConfig, loader *model.ModelLoader, responses chan schema.OpenAIResponse, extraUsage bool) error {
+	processTools := func(noAction string, prompt string, req *schema.OpenAIRequest, config *config.ModelConfig, loader *model.ModelLoader, responses chan schema.OpenAIResponse, extraUsage bool, id string, created int, textContentToReturn *string) error {
 		// Detect if thinking token is already in prompt or template
 		var template string
 		if config.TemplateConfig.UseTokenizerTemplate {
@@ -308,18 +305,18 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
 			xlog.Debug("[ChatDeltas] Using pre-parsed tool calls from C++ autoparser", "count", len(deltaToolCalls))
 			functionResults = deltaToolCalls
 			// Use content/reasoning from deltas too
-			textContentToReturn = functions.ContentFromChatDeltas(chatDeltas)
+			*textContentToReturn = functions.ContentFromChatDeltas(chatDeltas)
 			reasoning = functions.ReasoningFromChatDeltas(chatDeltas)
 		} else {
 			// Fallback: parse tool calls from raw text (no chat deltas from backend)
 			xlog.Debug("[ChatDeltas] no pre-parsed tool calls, falling back to Go-side text parsing")
 			reasoning = extractor.Reasoning()
 			cleanedResult := extractor.CleanedContent()
-			textContentToReturn = functions.ParseTextContent(cleanedResult, config.FunctionsConfig)
+			*textContentToReturn = functions.ParseTextContent(cleanedResult, config.FunctionsConfig)
 			cleanedResult = functions.CleanupLLMResult(cleanedResult, config.FunctionsConfig)
 			functionResults = functions.ParseFunctionCall(cleanedResult, config.FunctionsConfig)
 		}
-		xlog.Debug("[ChatDeltas] final tool call decision", "tool_calls", len(functionResults), "text_content", textContentToReturn)
+		xlog.Debug("[ChatDeltas] final tool call decision", "tool_calls", len(functionResults), "text_content", *textContentToReturn)
 		noActionToRun := len(functionResults) > 0 && functionResults[0].Name == noAction || len(functionResults) == 0
 
 		switch {
@@ -412,7 +409,7 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
 					Choices: []schema.Choice{{
 						Delta: &schema.Message{
 							Role:    "assistant",
-							Content: &textContentToReturn,
+							Content: textContentToReturn,
 							ToolCalls: []schema.ToolCall{
 								{
 									Index: i,
@@ -437,9 +434,9 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
 	}
 
 	return func(c echo.Context) error {
-		textContentToReturn = ""
-		id = uuid.New().String()
-		created = int(time.Now().Unix())
+		var textContentToReturn string
+		id := uuid.New().String()
+		created := int(time.Now().Unix())
 
 		input, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenAIRequest)
 		if !ok || input.Model == "" {
@@ -666,9 +663,9 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
 
 				go func() {
 					if !shouldUseFn {
-						ended <- process(predInput, input, config, ml, responses, extraUsage)
+						ended <- process(predInput, input, config, ml, responses, extraUsage, id, created)
 					} else {
-						ended <- processTools(noActionName, predInput, input, config, ml, responses, extraUsage)
+						ended <- processTools(noActionName, predInput, input, config, ml, responses, extraUsage, id, created, &textContentToReturn)
 					}
 				}()
 
@@ -747,6 +744,14 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
 					}
 				}
 
+				// Drain responses channel to unblock the background goroutine if it's
+				// still trying to send (e.g., after client disconnect). The goroutine
+				// calls close(responses) when done, which terminates the drain.
+				if input.Context.Err() != nil {
+					go func() { for range responses {} }()
+					<-ended
+				}
+
 				// MCP streaming tool execution: if we collected MCP tool calls, execute and loop
 				if hasMCPToolsStream && toolsCalled && len(collectedToolCalls) > 0 {
 					var hasMCPCalls bool

core/schema/anthropic.go

@@ -121,6 +121,7 @@ type AnthropicStreamEvent struct {
 	Delta        *AnthropicStreamDelta   `json:"delta,omitempty"`
 	Message      *AnthropicStreamMessage `json:"message,omitempty"`
 	Usage        *AnthropicUsage         `json:"usage,omitempty"`
+	Error        *AnthropicError         `json:"error,omitempty"`
 }
 
 // AnthropicStreamDelta represents the delta in a streaming response

core/services/advisorylock/advisorylock.go

@@ -36,80 +36,6 @@ func TryWithLockCtx(ctx context.Context, db *gorm.DB, key int64, fn func() error
 	return true, fn()
 }
 
-// Deprecated: Use TryWithLockCtx instead.
-//
-// TryWithLock attempts a non-blocking advisory lock on a dedicated connection.
-// If acquired, fn runs and the lock is released on the same connection.
-// Returns (true, fn-error) if acquired, (false, nil) if not.
-func TryWithLock(db *gorm.DB, key int64, fn func() error) (bool, error) {
-	sqlDB, err := db.DB()
-	if err != nil {
-		return false, fmt.Errorf("advisorylock: getting sql.DB: %w", err)
-	}
-	conn, err := sqlDB.Conn(context.Background())
-	if err != nil {
-		return false, fmt.Errorf("advisorylock: getting connection: %w", err)
-	}
-	defer conn.Close()
-
-	var acquired bool
-	if err := conn.QueryRowContext(context.Background(),
-		"SELECT pg_try_advisory_lock($1)", key).Scan(&acquired); err != nil {
-		return false, fmt.Errorf("advisorylock: try lock %d: %w", key, err)
-	}
-	if !acquired {
-		return false, nil
-	}
-	defer conn.ExecContext(context.Background(), "SELECT pg_advisory_unlock($1)", key)
-	return true, fn()
-}
-
-// Deprecated: TryLock may use different pool connections for lock/unlock.
-// Use TryWithLock instead, which pins a single connection.
-//
-// TryLock attempts to acquire a PostgreSQL advisory lock (non-blocking).
-// Returns true if the lock was acquired.
-func TryLock(db *gorm.DB, key int64) bool {
-	var acquired bool
-	db.Raw("SELECT pg_try_advisory_lock(?)", key).Scan(&acquired)
-	return acquired
-}
-
-// Deprecated: Unlock may use a different pool connection than TryLock.
-// Use TryWithLock instead, which pins a single connection.
-//
-// Unlock releases a PostgreSQL advisory lock.
-func Unlock(db *gorm.DB, key int64) {
-	db.Exec("SELECT pg_advisory_unlock(?)", key)
-}
-
-// Deprecated: Use WithLockCtx instead.
-//
-// WithLock acquires a PostgreSQL advisory lock for the duration of fn.
-// Uses a dedicated database connection to ensure the lock is held correctly.
-// The lock is automatically released when fn returns (or panics).
-func WithLock(db *gorm.DB, key int64, fn func() error) error {
-	sqlDB, err := db.DB()
-	if err != nil {
-		return fmt.Errorf("advisorylock: getting sql.DB: %w", err)
-	}
-
-	conn, err := sqlDB.Conn(context.Background())
-	if err != nil {
-		return fmt.Errorf("advisorylock: getting connection: %w", err)
-	}
-	defer conn.Close()
-
-	if _, err := conn.ExecContext(context.Background(),
-		"SELECT pg_advisory_lock($1)", key); err != nil {
-		return fmt.Errorf("advisorylock: acquiring lock %d: %w", key, err)
-	}
-	defer conn.ExecContext(context.Background(),
-		"SELECT pg_advisory_unlock($1)", key)
-
-	return fn()
-}
-
 // KeyFromUUID converts a UUID (as [16]byte) to a lock key via FNV-1a hashing.
 func KeyFromUUID(b [16]byte) int64 {
 	h := fnv.New64a()

core/services/advisorylock/advisorylock_test.go

@@ -40,7 +40,7 @@ var _ = Describe("AdvisoryLock", func() {
 			for range 2 {
 				wg.Go(func() {
 					defer GinkgoRecover()
-					err := WithLock(db, lockKey, func() error {
+					err := WithLockCtx(context.Background(), db, lockKey, func() error {
 						cur := atomic.AddInt32(&running, 1)
 						mu.Lock()
 						if cur > maxRunning {
@@ -66,18 +66,22 @@ var _ = Describe("AdvisoryLock", func() {
 			Expect(concurrency).To(BeZero(), "detected concurrent execution inside advisory lock")
 		})
 
-		It("acquires and unlocks with TryLock", func() {
+		It("acquires and releases with TryWithLockCtx", func() {
 			const lockKey int64 = 888
 
-			acquired := TryLock(db, lockKey)
-			Expect(acquired).To(BeTrue(), "expected TryLock to acquire the lock")
-
-			Unlock(db, lockKey)
-
-			reacquired := TryLock(db, lockKey)
-			Expect(reacquired).To(BeTrue(), "expected TryLock to re-acquire the lock after Unlock")
+			acquired, err := TryWithLockCtx(context.Background(), db, lockKey, func() error {
+				// Lock is held here
+				return nil
+			})
+			Expect(err).ToNot(HaveOccurred())
+			Expect(acquired).To(BeTrue(), "expected TryWithLockCtx to acquire the lock")
 
-			Unlock(db, lockKey)
+			// Lock released — should be re-acquirable
+			reacquired, err := TryWithLockCtx(context.Background(), db, lockKey, func() error {
+				return nil
+			})
+			Expect(err).ToNot(HaveOccurred())
+			Expect(reacquired).To(BeTrue(), "expected TryWithLockCtx to re-acquire the lock")
 		})
 	})
 

core/services/advisorylock/leader_loop.go

@@ -2,6 +2,7 @@ package advisorylock
 
 import (
 	"context"
+	"runtime/debug"
 	"time"
 
 	"github.com/mudler/xlog"
@@ -21,13 +22,20 @@ func RunLeaderLoop(ctx context.Context, db *gorm.DB, lockKey int64, interval tim
 		case <-ctx.Done():
 			return
 		case <-ticker.C:
-			_, err := TryWithLockCtx(ctx, db, lockKey, func() error {
-				fn()
-				return nil
-			})
-			if err != nil {
-				xlog.Error("Leader loop advisory lock error", "key", lockKey, "error", err)
-			}
+			func() {
+				defer func() {
+					if r := recover(); r != nil {
+						xlog.Error("Leader loop callback panicked", "key", lockKey, "panic", r, "stack", string(debug.Stack()))
+					}
+				}()
+				_, err := TryWithLockCtx(ctx, db, lockKey, func() error {
+					fn()
+					return nil
+				})
+				if err != nil {
+					xlog.Error("Leader loop advisory lock error", "key", lockKey, "error", err)
+				}
+			}()
 		}
 	}
 }