feat(distributed): add ExposeNodeHeader middleware + response writer wrapper

New ApplicationConfig.ExposeNodeHeader bool + --expose-node-header CLI
flag / LOCALAI_EXPOSE_NODE_HEADER env var (default off; the node ID
reveals internal topology and is opt-in).

The middleware creates a per-request *atomic.Value holder, attaches it
to c.Request().Context() via distributedhdr.WithHolder, and wraps
c.Response().Writer with a custom http.ResponseWriter that sets the
X-LocalAI-Node header on first Write / WriteHeader / Flush by reading
the holder. Implements http.Flusher, http.Hijacker, Unwrap so it
composes cleanly with Echo and http.NewResponseController.

request.go propagates the holder onto derived contexts via
distributedhdr.Inherit so the holder survives the correlation-ID
context replacement.

Unit + race-clean concurrency + integration specs.

Assisted-by: Claude:claude-opus-4-7[1m]
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

Author: mudler

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/middleware/node_header.go

@@ -0,0 +1,127 @@
+package middleware
+
+import (
+	"bufio"
+	"fmt"
+	"net"
+	"net/http"
+
+	"github.com/labstack/echo/v4"
+
+	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/pkg/distributedhdr"
+)
+
+// NodeHeaderName is the HTTP response header that, when --expose-node-header
+// is enabled, carries the ID of the distributed-mode worker node that served
+// the inference request. Off by default: node IDs reveal internal topology
+// and should not be exposed on a public endpoint.
+const NodeHeaderName = "X-LocalAI-Node"
+
+// nodeHeaderWriter wraps an http.ResponseWriter and stamps the X-LocalAI-Node
+// header lazily on the first Write / WriteHeader / Flush call. The lazy
+// resolve is what makes this work for streaming: the picked node ID is only
+// known AFTER the router runs (i.e. on the first SSE chunk), so resolving at
+// request entry would attach the previous request's routing decision (or
+// nothing on a cold cache).
+type nodeHeaderWriter struct {
+	http.ResponseWriter
+	resolve func() string
+	set     bool
+}
+
+func (w *nodeHeaderWriter) maybeSet() {
+	if w.set {
+		return
+	}
+	w.set = true
+	if id := w.resolve(); id != "" {
+		w.Header().Set(NodeHeaderName, id)
+	}
+}
+
+func (w *nodeHeaderWriter) Write(b []byte) (int, error) {
+	w.maybeSet()
+	return w.ResponseWriter.Write(b)
+}
+
+func (w *nodeHeaderWriter) WriteHeader(code int) {
+	w.maybeSet()
+	w.ResponseWriter.WriteHeader(code)
+}
+
+// Flush keeps SSE handlers working: Echo's Response.Flush goes through
+// http.NewResponseController which walks Unwrap() chains and invokes Flush
+// on the first wrapper that implements http.Flusher. By implementing it
+// here we both stamp the header before the underlying writer flushes AND
+// keep the streaming path alive.
+func (w *nodeHeaderWriter) Flush() {
+	w.maybeSet()
+	if f, ok := w.ResponseWriter.(http.Flusher); ok {
+		f.Flush()
+	}
+}
+
+// Hijack preserves WebSocket / raw-conn handlers that need to take over the
+// underlying TCP connection (e.g. /v1/realtime). Without this the wrapper
+// would silently break those endpoints.
+//
+// When the underlying writer does not implement http.Hijacker we return
+// http.ErrNotSupported so callers using errors.Is (notably
+// http.NewResponseController.Hijack) detect the condition through the
+// standard sentinel rather than a string-matched custom error.
+func (w *nodeHeaderWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if h, ok := w.ResponseWriter.(http.Hijacker); ok {
+		return h.Hijack()
+	}
+	return nil, nil, fmt.Errorf("hijack not supported: %w", http.ErrNotSupported)
+}
+
+// Unwrap lets http.NewResponseController reach through us to find optional
+// interfaces (CloseNotifier, SetReadDeadline, etc.) on the real writer.
+func (w *nodeHeaderWriter) Unwrap() http.ResponseWriter {
+	return w.ResponseWriter
+}
+
+// ExposeNodeHeader installs a per-request response writer wrapper that
+// stamps the X-LocalAI-Node header from the per-request holder published
+// by the distributed router on the first write. Off by default; opted in
+// via --expose-node-header / LOCALAI_EXPOSE_NODE_HEADER.
+//
+// Attribution is per-request correct: the middleware creates a fresh
+// holder per request, plumbs it through context.Context, and the router
+// writes the picked node ID for THIS request's routing decision. No
+// shared loader state, no overwriting across concurrent requests for the
+// same model on multiple replicas.
+func ExposeNodeHeader(appCfg *config.ApplicationConfig) echo.MiddlewareFunc {
+	return func(next echo.HandlerFunc) echo.HandlerFunc {
+		return func(c echo.Context) error {
+			if appCfg == nil || !appCfg.ExposeNodeHeader {
+				return next(c)
+			}
+
+			// One holder per request. The pointer is captured both in
+			// the wrapper closure (read side) and in the request
+			// context (write side, accessed by the router via
+			// distributedhdr.Stamp). Both sides point at the same
+			// atomic slot.
+			holder := distributedhdr.NewHolder()
+
+			req := c.Request()
+			c.SetRequest(req.WithContext(distributedhdr.WithHolder(req.Context(), holder)))
+
+			orig := c.Response().Writer
+			wrapper := &nodeHeaderWriter{
+				ResponseWriter: orig,
+				resolve: func() string {
+					return distributedhdr.Load(holder)
+				},
+			}
+			c.Response().Writer = wrapper
+			defer func() {
+				c.Response().Writer = orig
+			}()
+			return next(c)
+		}
+	}
+}

core/http/middleware/node_header_concurrency_test.go

@@ -0,0 +1,89 @@
+package middleware
+
+// Regression coverage for the multi-replica X-LocalAI-Node attribution
+// bug fixed by this PR.
+//
+// Pre-refactor failure mode: ExposeNodeHeader resolved the node ID by
+// calling ml.LookupNodeID(modelName), which read a single per-modelID
+// slot in the loader's in-memory store. The distributed router
+// overwrote that slot on every routing decision, so when N concurrent
+// requests for the same model were routed to N different replicas, the
+// header value the wrapper picked up at first-byte time depended on
+// goroutine interleaving and not on which replica THIS request was
+// actually sent to.
+//
+// The fix routes attribution through a per-request atomic holder
+// installed by the middleware and stamped by the router via
+// distributedhdr.Stamp. Each request carries its own slot, so peer
+// stamps cannot bleed in.
+//
+// This spec exercises the exact concurrency pattern the bug required to
+// reproduce: many goroutines, all running through the same middleware
+// instance (mirroring e.Use() in production), each stamping a distinct
+// node ID, each asserting its own response carries the matching ID.
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+
+	"github.com/labstack/echo/v4"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/pkg/distributedhdr"
+)
+
+var _ = Describe("ExposeNodeHeader multi-replica attribution", func() {
+	It("each concurrent request sees the node ID stamped on ITS OWN request, not a peer's", func() {
+		appCfg := &config.ApplicationConfig{ExposeNodeHeader: true}
+		e := echo.New()
+		mw := ExposeNodeHeader(appCfg)
+
+		const N = 64
+		results := make([]string, N)
+		var wg sync.WaitGroup
+		wg.Add(N)
+
+		// Drive N concurrent requests through the same middleware
+		// instance. Each handler stamps a distinct, per-request node ID
+		// derived from the request index, then yields before writing
+		// the body so the goroutines have ample opportunity to
+		// interleave. Under the old shared-loader design this is the
+		// configuration that surfaced the bug; under the new
+		// per-request-holder design every request must round-trip its
+		// own stamp.
+		for i := 0; i < N; i++ {
+			i := i
+			go func() {
+				defer wg.Done()
+				req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
+				rec := httptest.NewRecorder()
+				c := e.NewContext(req, rec)
+
+				expected := fmt.Sprintf("node-%d", i)
+				handler := func(c echo.Context) error {
+					distributedhdr.Stamp(c.Request().Context(), expected)
+					// Yield to amplify interleaving: if any stamp were
+					// shared across requests, the late writes would
+					// observe peer-stamped values instead of their
+					// own.
+					for j := 0; j < 16; j++ {
+						_ = j
+					}
+					return c.String(http.StatusOK, "ok")
+				}
+				Expect(mw(handler)(c)).To(Succeed())
+				results[i] = rec.Header().Get(NodeHeaderName)
+			}()
+		}
+		wg.Wait()
+
+		for i := 0; i < N; i++ {
+			Expect(results[i]).To(Equal(fmt.Sprintf("node-%d", i)),
+				"request %d must see ITS OWN routing decision in the header, not a peer's", i)
+		}
+	})
+})