feat(distributed): gated X-LocalAI-Node response header (middleware + wrapper)

core/backend/ctx_propagation_test.go

@@ -0,0 +1,169 @@
+package backend_test
+
+// Regression spec for X-LocalAI-Node coverage on audio/image/TTS/rerank/VAD.
+//
+// The X-LocalAI-Node middleware (core/http/middleware.ExposeNodeHeader)
+// works end-to-end only if the per-request holder attached to the HTTP
+// request context reaches the SmartRouter via ml.Load(opts...). The chain
+// is:
+//
+//   handler -> backend.Foo(ctx, ...) -> ModelOptions(cfg, app, WithContext(ctx))
+//     -> ml.Load(opts...) -> grpcModel(..., o.context) -> modelRouter(ctx, ...)
+//     -> SmartRouter -> distributedhdr.Stamp(ctx, nodeID)
+//
+// If any backend helper drops `ctx` and lets ModelOptions fall back to the
+// app context, the router never sees the per-request holder and the
+// header silently stays empty for that endpoint. These specs pin the
+// request-context-reaches-router contract for the five backend helpers
+// that were previously dropping ctx between the handler and Load.
+
+import (
+	"context"
+	"sync/atomic"
+
+	"github.com/mudler/LocalAI/core/backend"
+	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/core/schema"
+	pbproto "github.com/mudler/LocalAI/pkg/grpc/proto"
+	"github.com/mudler/LocalAI/pkg/distributedhdr"
+	"github.com/mudler/LocalAI/pkg/model"
+	"github.com/mudler/LocalAI/pkg/system"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+// newCapturingLoader returns a ModelLoader wired with a stub model router
+// that captures the context it receives and then short-circuits with a
+// sentinel error. The router callback is the exact seam where the
+// SmartRouter would call distributedhdr.Stamp in production, so observing
+// the holder here is equivalent to observing it at the real router.
+func newCapturingLoader() (*model.ModelLoader, *atomic.Value, func() context.Context) {
+	loader := model.NewModelLoader(&system.SystemState{})
+	var captured atomic.Value
+	loader.SetModelRouter(func(ctx context.Context, _ string, _, _, _ string, _ *pbproto.ModelOptions, _ bool) (*model.Model, error) {
+		captured.Store(ctx)
+		// Return an error so the backend short-circuits before trying to
+		// dial gRPC. We only care about the context-arrival contract.
+		return nil, errRouterShortCircuit
+	})
+	get := func() context.Context {
+		v, _ := captured.Load().(context.Context)
+		return v
+	}
+	return loader, &captured, get
+}
+
+var errRouterShortCircuit = sentinelErr("router short-circuit (test)")
+
+type sentinelErr string
+
+func (s sentinelErr) Error() string { return string(s) }
+
+func newAppCfg() *config.ApplicationConfig {
+	return config.NewApplicationConfig(config.WithSystemState(&system.SystemState{}))
+}
+
+func newModelCfg() config.ModelConfig {
+	threads := 1
+	cfg := config.ModelConfig{
+		Name:    "test-model",
+		Backend: "stub-backend",
+		Threads: &threads,
+	}
+	cfg.Model = "test.bin"
+	return cfg
+}
+
+var _ = Describe("X-LocalAI-Node ctx propagation contract", func() {
+	const fakeNodeID = "node-ctx-propagation-7"
+
+	var (
+		appCfg      *config.ApplicationConfig
+		modelCfg    config.ModelConfig
+		loader      *model.ModelLoader
+		routerCtxOf func() context.Context
+		holder      *atomic.Value
+		reqCtx      context.Context
+	)
+
+	BeforeEach(func() {
+		appCfg = newAppCfg()
+		modelCfg = newModelCfg()
+		loader, _, routerCtxOf = newCapturingLoader()
+		holder = distributedhdr.NewHolder()
+		reqCtx = distributedhdr.WithHolder(context.Background(), holder)
+	})
+
+	// stampViaRouterCtx asserts the captured router context carries the
+	// SAME holder that was attached to the request. We verify by stamping
+	// through the router-side ctx and observing the value via the
+	// request-side holder; if the holders were different objects the load
+	// would return "".
+	stampViaRouterCtx := func() {
+		routerCtx := routerCtxOf()
+		Expect(routerCtx).ToNot(BeNil(), "router callback must have been invoked")
+		distributedhdr.Stamp(routerCtx, fakeNodeID)
+		Expect(distributedhdr.Load(holder)).To(Equal(fakeNodeID),
+			"stamp via router-side ctx must be observable via the request-side holder")
+	}
+
+	It("Rerank forwards the request context to the SmartRouter", func() {
+		_, err := backend.Rerank(reqCtx, &pbproto.RerankRequest{Query: "q"}, loader, appCfg, modelCfg)
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("router short-circuit (test)"))
+		stampViaRouterCtx()
+	})
+
+	It("VAD forwards the request context to the SmartRouter", func() {
+		_, err := backend.VAD(&schema.VADRequest{}, reqCtx, loader, appCfg, modelCfg)
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("router short-circuit (test)"))
+		stampViaRouterCtx()
+	})
+
+	It("ModelTTS forwards the request context to the SmartRouter", func() {
+		_, _, err := backend.ModelTTS(reqCtx, "hello", "", "", loader, appCfg, modelCfg)
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("router short-circuit (test)"))
+		stampViaRouterCtx()
+	})
+
+	It("ModelTTSStream forwards the request context to the SmartRouter", func() {
+		err := backend.ModelTTSStream(reqCtx, "hello", "", "", loader, appCfg, modelCfg, func([]byte) error { return nil })
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("router short-circuit (test)"))
+		stampViaRouterCtx()
+	})
+
+	It("ModelTranscriptionWithOptions forwards the request context to the SmartRouter", func() {
+		_, err := backend.ModelTranscriptionWithOptions(reqCtx, backend.TranscriptionRequest{Audio: "x.wav"}, loader, modelCfg, appCfg)
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("router short-circuit (test)"))
+		stampViaRouterCtx()
+	})
+
+	It("ModelTranscriptionStream forwards the request context to the SmartRouter", func() {
+		err := backend.ModelTranscriptionStream(reqCtx, backend.TranscriptionRequest{Audio: "x.wav"}, loader, modelCfg, appCfg, func(backend.TranscriptionStreamChunk) {})
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("router short-circuit (test)"))
+		stampViaRouterCtx()
+	})
+
+	It("ImageGeneration forwards the request context to the SmartRouter", func() {
+		_, err := backend.ImageGeneration(reqCtx, 64, 64, 1, 0, "p", "", "", "/tmp/out.png", loader, modelCfg, appCfg, nil)
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("router short-circuit (test)"))
+		stampViaRouterCtx()
+	})
+
+	It("does NOT leak the holder when the app context is used instead", func() {
+		// Sanity: the bug being fixed manifests as the router getting
+		// appCfg.Context (no holder) instead of reqCtx (holder). A direct
+		// call with context.Background() must not see the holder via the
+		// app context surface.
+		appCtxOnly := appCfg.Context
+		Expect(distributedhdr.Holder(appCtxOnly)).To(BeNil(),
+			"the app context must not be the carrier of per-request holders")
+	})
+})

core/backend/embeddings.go

@@ -30,17 +30,20 @@ type modelEmbedder struct {
 	appConfig   *config.ApplicationConfig
 }
 
-func (e *modelEmbedder) Embed(_ context.Context, text string) ([]float32, error) {
-	fn, err := ModelEmbedding(text, nil, e.loader, e.modelConfig, e.appConfig)
+func (e *modelEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
+	fn, err := ModelEmbedding(ctx, text, nil, e.loader, e.modelConfig, e.appConfig)
 	if err != nil {
 		return nil, err
 	}
 	return fn()
 }
 
-func ModelEmbedding(s string, tokens []int, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (func() ([]float32, error), error) {
+func ModelEmbedding(ctx context.Context, s string, tokens []int, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (func() ([]float32, error), error) {
 
-	opts := ModelOptions(modelConfig, appConfig)
+	// model.WithContext(ctx) overrides the app-context default set in
+	// ModelOptions so distributed routing decisions reach the request's
+	// X-LocalAI-Node holder via distributedhdr.Stamp.
+	opts := ModelOptions(modelConfig, appConfig, model.WithContext(ctx))
 
 	inferenceModel, err := loader.Load(opts...)
 	if err != nil {

core/backend/image.go

@@ -1,6 +1,7 @@
 package backend
 
 import (
+	"context"
 	"time"
 
 	"github.com/mudler/LocalAI/core/config"
@@ -10,9 +11,12 @@ import (
 	model "github.com/mudler/LocalAI/pkg/model"
 )
 
-func ImageGeneration(height, width, step, seed int, positive_prompt, negative_prompt, src, dst string, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig, refImages []string) (func() error, error) {
+func ImageGeneration(ctx context.Context, height, width, step, seed int, positive_prompt, negative_prompt, src, dst string, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig, refImages []string) (func() error, error) {
 
-	opts := ModelOptions(modelConfig, appConfig)
+	// model.WithContext(ctx) overrides the app-context default set in
+	// ModelOptions so distributed routing decisions reach the request's
+	// X-LocalAI-Node holder via distributedhdr.Stamp.
+	opts := ModelOptions(modelConfig, appConfig, model.WithContext(ctx))
 	inferenceModel, err := loader.Load(
 		opts...,
 	)
@@ -23,7 +27,7 @@ func ImageGeneration(height, width, step, seed int, positive_prompt, negative_pr
 
 	fn := func() error {
 		_, err := inferenceModel.GenerateImage(
-			appConfig.Context,
+			ctx,
 			&proto.GenerateImageRequest{
 				Height:           int32(height),
 				Width:            int32(width),

core/backend/llm.go

@@ -94,7 +94,7 @@ func ModelInference(ctx context.Context, s string, messages schema.Messages, ima
 		}
 	}
 
-	opts := ModelOptions(*c, o)
+	opts := ModelOptions(*c, o, model.WithContext(ctx))
 	inferenceModel, err := loader.Load(opts...)
 	if err != nil {
 		recordModelLoadFailure(o, c.Name, c.Backend, err, map[string]any{"model_file": modelFile})

core/backend/rerank.go

@@ -20,7 +20,7 @@ type RerankResult struct {
 }
 
 // Reranker scores a list of candidate documents against a query.
-// Returns one RerankResult per input document (no top-N truncation —
+// Returns one RerankResult per input document (no top-N truncation -
 // callers that need it can sort and slice).
 type Reranker interface {
 	Rerank(ctx context.Context, query string, documents []string) ([]RerankResult, error)
@@ -41,7 +41,7 @@ func (r *modelReranker) Rerank(ctx context.Context, query string, documents []st
 	req := &proto.RerankRequest{
 		Query:     query,
 		Documents: documents,
-		// TopN=0 → backend returns scores for every document. Truncating
+		// TopN=0: backend returns scores for every document. Truncating
 		// here would silently zero out labels the reranker considered
 		// unlikely, which the router classifier needs.
 	}
@@ -57,7 +57,10 @@ func (r *modelReranker) Rerank(ctx context.Context, query string, documents []st
 }
 
 func Rerank(ctx context.Context, request *proto.RerankRequest, loader *model.ModelLoader, appConfig *config.ApplicationConfig, modelConfig config.ModelConfig) (*proto.RerankResult, error) {
-	opts := ModelOptions(modelConfig, appConfig)
+	// model.WithContext(ctx) overrides the app-context default set in
+	// ModelOptions so distributed routing decisions reach the request's
+	// X-LocalAI-Node holder via distributedhdr.Stamp.
+	opts := ModelOptions(modelConfig, appConfig, model.WithContext(ctx))
 	rerankModel, err := loader.Load(opts...)
 	if err != nil {
 		recordModelLoadFailure(appConfig, modelConfig.Name, modelConfig.Backend, err, nil)

core/backend/transcript.go

@@ -41,11 +41,14 @@ func (r *TranscriptionRequest) toProto(threads uint32) *proto.TranscriptRequest
 	}
 }
 
-func loadTranscriptionModel(ml *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (grpcPkg.Backend, error) {
+func loadTranscriptionModel(ctx context.Context, ml *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (grpcPkg.Backend, error) {
 	if modelConfig.Backend == "" {
 		modelConfig.Backend = model.WhisperBackend
 	}
-	opts := ModelOptions(modelConfig, appConfig)
+	// model.WithContext(ctx) overrides the app-context default set in
+	// ModelOptions so distributed routing decisions reach the request's
+	// X-LocalAI-Node holder via distributedhdr.Stamp.
+	opts := ModelOptions(modelConfig, appConfig, model.WithContext(ctx))
 	transcriptionModel, err := ml.Load(opts...)
 	if err != nil {
 		recordModelLoadFailure(appConfig, modelConfig.Name, modelConfig.Backend, err, nil)
@@ -68,7 +71,7 @@ func ModelTranscription(ctx context.Context, audio, language string, translate,
 }
 
 func ModelTranscriptionWithOptions(ctx context.Context, req TranscriptionRequest, ml *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (*schema.TranscriptionResult, error) {
-	transcriptionModel, err := loadTranscriptionModel(ml, modelConfig, appConfig)
+	transcriptionModel, err := loadTranscriptionModel(ctx, ml, modelConfig, appConfig)
 	if err != nil {
 		return nil, err
 	}
@@ -150,7 +153,7 @@ type TranscriptionStreamChunk struct {
 // support real streaming should still emit one terminal event with Final set,
 // which the HTTP layer turns into a single delta + done SSE pair.
 func ModelTranscriptionStream(ctx context.Context, req TranscriptionRequest, ml *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig, onChunk func(TranscriptionStreamChunk)) error {
-	transcriptionModel, err := loadTranscriptionModel(ml, modelConfig, appConfig)
+	transcriptionModel, err := loadTranscriptionModel(ctx, ml, modelConfig, appConfig)
 	if err != nil {
 		return err
 	}

core/backend/tts.go

@@ -29,7 +29,10 @@ func ModelTTS(
 	appConfig *config.ApplicationConfig,
 	modelConfig config.ModelConfig,
 ) (string, *proto.Result, error) {
-	opts := ModelOptions(modelConfig, appConfig)
+	// model.WithContext(ctx) overrides the app-context default set in
+	// ModelOptions so distributed routing decisions reach the request's
+	// X-LocalAI-Node holder via distributedhdr.Stamp.
+	opts := ModelOptions(modelConfig, appConfig, model.WithContext(ctx))
 	ttsModel, err := loader.Load(opts...)
 	if err != nil {
 		recordModelLoadFailure(appConfig, modelConfig.Name, modelConfig.Backend, err, nil)
@@ -131,7 +134,7 @@ func ModelTTSStream(
 	modelConfig config.ModelConfig,
 	audioCallback func([]byte) error,
 ) error {
-	opts := ModelOptions(modelConfig, appConfig)
+	opts := ModelOptions(modelConfig, appConfig, model.WithContext(ctx))
 	ttsModel, err := loader.Load(opts...)
 	if err != nil {
 		recordModelLoadFailure(appConfig, modelConfig.Name, modelConfig.Backend, err, nil)

core/backend/vad.go

@@ -14,7 +14,10 @@ func VAD(request *schema.VADRequest,
 	ml *model.ModelLoader,
 	appConfig *config.ApplicationConfig,
 	modelConfig config.ModelConfig) (*schema.VADResponse, error) {
-	opts := ModelOptions(modelConfig, appConfig)
+	// model.WithContext(ctx) overrides the app-context default set in
+	// ModelOptions so distributed routing decisions reach the request's
+	// X-LocalAI-Node holder via distributedhdr.Stamp.
+	opts := ModelOptions(modelConfig, appConfig, model.WithContext(ctx))
 	vadModel, err := ml.Load(opts...)
 	if err != nil {
 		recordModelLoadFailure(appConfig, modelConfig.Name, modelConfig.Backend, err, nil)

core/cli/run.go

@@ -157,6 +157,7 @@ type RunCMD struct {
 	AutoApproveNodes      bool   `env:"LOCALAI_AUTO_APPROVE_NODES" default:"false" help:"Auto-approve new worker nodes (skip admin approval)" group:"distributed"`
 	BackendInstallTimeout string `env:"LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT" help:"NATS round-trip timeout for backend.install requests sent to worker nodes (default 15m). Increase for slow links pulling multi-GB images." group:"distributed"`
 	BackendUpgradeTimeout string `env:"LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT" help:"NATS round-trip timeout for backend.upgrade requests (default 15m)." group:"distributed"`
+	ExposeNodeHeader      bool   `env:"LOCALAI_EXPOSE_NODE_HEADER" default:"false" help:"Set the X-LocalAI-Node response header on inference responses (OpenAI chat/completions/embeddings, Anthropic /v1/messages, Ollama /api/chat,/api/generate,/api/embed) with the ID of the worker that served the request. Disabled by default: the node ID reveals internal topology and should not be exposed on a public endpoint. Best-effort: under heavy concurrency the header may reflect a recent routing decision rather than this exact request's." group:"distributed"`
 
 	Version bool
 
@@ -283,6 +284,9 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
 	if r.AutoApproveNodes {
 		opts = append(opts, config.EnableAutoApproveNodes)
 	}
+	if r.ExposeNodeHeader {
+		opts = append(opts, config.WithExposeNodeHeader(true))
+	}
 
 	if r.DisableMetricsEndpoint {
 		opts = append(opts, config.DisableMetricsEndpoint)

core/config/application_config.go

@@ -160,6 +160,18 @@ type ApplicationConfig struct {
 	// Distributed / Horizontal Scaling
 	Distributed DistributedConfig
 
+	// ExposeNodeHeader, when true, activates middleware.ExposeNodeHeader on
+	// the inference routes (OpenAI chat/completions/embeddings, Anthropic
+	// /v1/messages, Ollama /api/chat,/api/generate,/api/embed). The
+	// middleware wraps the response writer and attaches an "X-LocalAI-Node"
+	// response header carrying the ID of the distributed-mode worker node
+	// that served the request. Off by default because the node ID is
+	// internal topology that can aid attacker reconnaissance if surfaced on
+	// a public endpoint; operators opt in explicitly via
+	// --expose-node-header / LOCALAI_EXPOSE_NODE_HEADER for debugging,
+	// observability and load-balancer attribution.
+	ExposeNodeHeader bool
+
 	// LocalAI Assistant chat modality. Hard-disable the in-process admin MCP
 	// server with this flag; runtime-toggleable via /api/settings.
 	DisableLocalAIAssistant bool
@@ -980,6 +992,15 @@ func WithDisableLocalAIAssistant(disabled bool) AppOption {
 	}
 }
 
+// WithExposeNodeHeader enables the X-LocalAI-Node response header on
+// inference endpoints. Default off; the node ID reveals internal cluster
+// topology and is opt-in for that reason.
+func WithExposeNodeHeader(enabled bool) AppOption {
+	return func(o *ApplicationConfig) {
+		o.ExposeNodeHeader = enabled
+	}
+}
+
 // ToConfigLoaderOptions returns a slice of ConfigLoader Option.
 // Some options defined at the application level are going to be passed as defaults for
 // all the configuration for the models.

core/http/endpoints/ollama/embed.go

@@ -36,7 +36,7 @@ func EmbedEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfi
 		promptEvalCount := 0
 
 		for _, s := range inputStrings {
-			embedFn, err := backend.ModelEmbedding(s, []int{}, ml, *cfg, appConfig)
+			embedFn, err := backend.ModelEmbedding(c.Request().Context(), s, []int{}, ml, *cfg, appConfig)
 			if err != nil {
 				xlog.Error("Ollama embed failed", "error", err)
 				return ollamaError(c, 500, fmt.Sprintf("embedding failed: %v", err))