From 06e777b75e3201a0b7ab4b065624b79d0bba9012 Mon Sep 17 00:00:00 2001 From: "LocalAI [bot]" <139863280+localai-bot@users.noreply.github.com> Date: Mon, 25 May 2026 10:51:48 +0200 Subject: [PATCH] feat(distributed): gated X-LocalAI-Node response header (middleware + wrapper) (#9976) * feat(distributed): add per-request node ID context holder Introduce pkg/distributedhdr, a leaf package carrying a per-request *atomic.Value holder for the picked worker node ID from the SmartRouter (core/services/nodes) up to the HTTP response writer wrapper (core/http/middleware). Avoids the import cycle that a shared key in either consumer would create. Exposes NewHolder, WithHolder, Holder, Stamp, Load, Inherit. The holder is atomic.Value so cross-goroutine publish from the router to the response writer wrapper is race-clean. Assisted-by: Claude:claude-opus-4-7[1m] Signed-off-by: Ettore Di Giacinto * 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 * feat(distributed): stamp node ID in router and wire middleware to inference routes ModelRouterAdapter.Route stamps the picked node ID into the per-request holder via distributedhdr.Stamp(ctx, result.Node.ID) right after replica selection. Wire ExposeNodeHeader middleware to: - OpenAI chat/completion/embeddings + audio transcriptions/speech + image generations/inpainting - Anthropic /v1/messages - Ollama /api/chat, /api/generate, /api/embed, /api/embeddings - Jina /v1/rerank - LocalAI /v1/vad The middleware's wrapper reads the holder on first byte and sets the X-LocalAI-Node response header before delegating to the underlying writer. Per-request scope means no race under concurrent multi-replica routing. Assisted-by: Claude:claude-opus-4-7[1m] Signed-off-by: Ettore Di Giacinto * fix(distributed): thread request context through backend Load + cover ctx propagation Five non-OpenAI backend helpers were silently using app.Context instead of the request context for the gRPC backend call: transcription, TTS, image generation, rerank, VAD. Effect: distributedhdr.Stamp in the router callback was a silent no-op for these paths, AND client cancellation didn't propagate to in-flight inference. Thread c.Request().Context() (or the equivalent input.Context after the request middleware has installed the correlation-ID derived context) through each helper and into ModelOptions via model.WithContext(ctx). ImageGeneration's signature gains a leading ctx parameter; in-tree callers (openai image, openai inpainting, openai inpainting_test) are updated to match. ModelEmbedding gains a leading ctx parameter for the same reason; the openai and ollama embedding handlers pass the request context through. chat_stream_workers.go defers the initial role=assistant chunk emission until the first token callback so the wrapper's lazy X-LocalAI-Node lookup against the loader runs AFTER ml.Load has stamped the per-modelID node ID; semantically identical for clients (role still arrives before any text). Regression test core/backend/ctx_propagation_test.go pins ctx propagation for all five helpers. Docs updated to enumerate the full endpoint coverage of the --expose-node-header flag. Assisted-by: Claude:claude-opus-4-7[1m] Signed-off-by: Ettore Di Giacinto --------- Signed-off-by: Ettore Di Giacinto Co-authored-by: Ettore Di Giacinto --- core/backend/ctx_propagation_test.go | 169 ++++++++++++ core/backend/embeddings.go | 11 +- core/backend/image.go | 10 +- core/backend/llm.go | 2 +- core/backend/rerank.go | 9 +- core/backend/transcript.go | 11 +- core/backend/tts.go | 7 +- core/backend/vad.go | 5 +- core/cli/run.go | 4 + core/config/application_config.go | 21 ++ core/http/endpoints/ollama/embed.go | 2 +- .../endpoints/openai/chat_stream_workers.go | 29 +- core/http/endpoints/openai/embeddings.go | 4 +- core/http/endpoints/openai/image.go | 2 +- core/http/endpoints/openai/inpainting.go | 2 +- core/http/endpoints/openai/inpainting_test.go | 3 +- core/http/middleware/node_header.go | 127 +++++++++ .../node_header_concurrency_test.go | 89 ++++++ .../node_header_integration_test.go | 201 ++++++++++++++ core/http/middleware/node_header_test.go | 260 ++++++++++++++++++ core/http/middleware/request.go | 3 + core/http/routes/anthropic.go | 3 + core/http/routes/jina.go | 1 + core/http/routes/localai.go | 3 + core/http/routes/ollama.go | 7 + core/http/routes/openai.go | 13 + core/services/nodes/model_router.go | 11 + docs/content/features/distributed-mode.md | 1 + pkg/distributedhdr/distributedhdr.go | 108 ++++++++ pkg/distributedhdr/distributedhdr_test.go | 99 +++++++ 30 files changed, 1186 insertions(+), 31 deletions(-) create mode 100644 core/backend/ctx_propagation_test.go create mode 100644 core/http/middleware/node_header.go create mode 100644 core/http/middleware/node_header_concurrency_test.go create mode 100644 core/http/middleware/node_header_integration_test.go create mode 100644 core/http/middleware/node_header_test.go create mode 100644 pkg/distributedhdr/distributedhdr.go create mode 100644 pkg/distributedhdr/distributedhdr_test.go diff --git a/core/backend/ctx_propagation_test.go b/core/backend/ctx_propagation_test.go new file mode 100644 index 000000000..f95549ce9 --- /dev/null +++ b/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") + }) +}) diff --git a/core/backend/embeddings.go b/core/backend/embeddings.go index dd9b9cffb..4be2bc346 100644 --- a/core/backend/embeddings.go +++ b/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 { diff --git a/core/backend/image.go b/core/backend/image.go index cd3b8ce6b..7414d5d2b 100644 --- a/core/backend/image.go +++ b/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), diff --git a/core/backend/llm.go b/core/backend/llm.go index 572d943e0..053e984e8 100644 --- a/core/backend/llm.go +++ b/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}) diff --git a/core/backend/rerank.go b/core/backend/rerank.go index f82208f8c..0c0d15b22 100644 --- a/core/backend/rerank.go +++ b/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) diff --git a/core/backend/transcript.go b/core/backend/transcript.go index e9b5f5360..e6da923cc 100644 --- a/core/backend/transcript.go +++ b/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 } diff --git a/core/backend/tts.go b/core/backend/tts.go index 62c394714..3ddc159d8 100644 --- a/core/backend/tts.go +++ b/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) diff --git a/core/backend/vad.go b/core/backend/vad.go index c3ecb66c9..1c874e338 100644 --- a/core/backend/vad.go +++ b/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) diff --git a/core/cli/run.go b/core/cli/run.go index 0f01e2ab4..725fc0939 100644 --- a/core/cli/run.go +++ b/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) diff --git a/core/config/application_config.go b/core/config/application_config.go index 24e7a82fc..dd36b97b9 100644 --- a/core/config/application_config.go +++ b/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. diff --git a/core/http/endpoints/ollama/embed.go b/core/http/endpoints/ollama/embed.go index 304343cf6..2d7428b5b 100644 --- a/core/http/endpoints/ollama/embed.go +++ b/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)) diff --git a/core/http/endpoints/openai/chat_stream_workers.go b/core/http/endpoints/openai/chat_stream_workers.go index f06da9111..e5eb2d1ca 100644 --- a/core/http/endpoints/openai/chat_stream_workers.go +++ b/core/http/endpoints/openai/chat_stream_workers.go @@ -21,6 +21,13 @@ import ( // The caller owns the `responses` channel and is expected to read from // it while this function runs; processStream closes the channel before // returning. +// +// X-LocalAI-Node attribution (when --expose-node-header is on) is +// handled by middleware.ExposeNodeHeader at the response writer wrapper +// layer; no in-band signal from the worker is needed. The initial +// role=assistant chunk is still emitted from the first token callback +// rather than eagerly here, so the wrapper's lazy lookup against the +// loader runs AFTER ml.Load has stamped the per-modelID node ID. func processStream( s string, req *schema.OpenAIRequest, @@ -32,13 +39,7 @@ func processStream( id string, created int, ) (backend.TokenUsage, error) { - responses <- schema.OpenAIResponse{ - ID: id, - Created: created, - Model: req.Model, // we have to return what the user sent here, due to OpenAI spec. - Choices: []schema.Choice{{Delta: &schema.Message{Role: "assistant"}, Index: 0, FinishReason: nil}}, - Object: "chat.completion.chunk", - } + sentInitialRole := false // Detect if thinking token is already in prompt or template // When UseTokenizerTemplate is enabled, predInput is empty, so we check the template @@ -70,6 +71,17 @@ func processStream( contentDelta = goContent } + if !sentInitialRole { + sentInitialRole = true + responses <- schema.OpenAIResponse{ + ID: id, + Created: created, + Model: req.Model, // we have to return what the user sent here, due to OpenAI spec. + Choices: []schema.Choice{{Delta: &schema.Message{Role: "assistant"}, Index: 0, FinishReason: nil}}, + Object: "chat.completion.chunk", + } + } + delta := &schema.Message{} if contentDelta != "" { delta.Content = &contentDelta @@ -130,6 +142,9 @@ func processStreamWithTools( hasChatDeltaToolCalls := false hasChatDeltaContent := false + // X-LocalAI-Node attribution is handled by middleware.ExposeNodeHeader + // at the wrapper layer; no in-band signalling from this worker. + _, finalUsage, chatDeltas, err := ComputeChoices(req, prompt, cfg, cl, startupOptions, loader, func(s string, c *[]schema.Choice) {}, func(s string, usage backend.TokenUsage) bool { result += s diff --git a/core/http/endpoints/openai/embeddings.go b/core/http/endpoints/openai/embeddings.go index 96fd5efc8..e6a3a353d 100644 --- a/core/http/endpoints/openai/embeddings.go +++ b/core/http/endpoints/openai/embeddings.go @@ -63,7 +63,7 @@ func EmbeddingsEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, app for i, s := range config.InputToken { // get the model function to call for the result - embedFn, err := backend.ModelEmbedding("", s, ml, *config, appConfig) + embedFn, err := backend.ModelEmbedding(input.Context, "", s, ml, *config, appConfig) if err != nil { return err } @@ -77,7 +77,7 @@ func EmbeddingsEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, app for i, s := range config.InputStrings { // get the model function to call for the result - embedFn, err := backend.ModelEmbedding(s, []int{}, ml, *config, appConfig) + embedFn, err := backend.ModelEmbedding(input.Context, s, []int{}, ml, *config, appConfig) if err != nil { return err } diff --git a/core/http/endpoints/openai/image.go b/core/http/endpoints/openai/image.go index 47e53225f..b083a2dbf 100644 --- a/core/http/endpoints/openai/image.go +++ b/core/http/endpoints/openai/image.go @@ -198,7 +198,7 @@ func ImageEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfi inputSrc = inputImages[0] } - fn, err := backend.ImageGeneration(height, width, step, *config.Seed, positive_prompt, negative_prompt, inputSrc, output, ml, *config, appConfig, refImages) + fn, err := backend.ImageGeneration(c.Request().Context(), height, width, step, *config.Seed, positive_prompt, negative_prompt, inputSrc, output, ml, *config, appConfig, refImages) if err != nil { return err } diff --git a/core/http/endpoints/openai/inpainting.go b/core/http/endpoints/openai/inpainting.go index 5cf37d34e..d0d0af585 100644 --- a/core/http/endpoints/openai/inpainting.go +++ b/core/http/endpoints/openai/inpainting.go @@ -231,7 +231,7 @@ func InpaintingEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, app // Note: ImageGenerationFunc will call into the loaded model's GenerateImage which expects src JSON // Also pass ref images (orig + mask) so backends that support ref images can use them. refImages := []string{origRef, maskRef} - fn, err := backend.ImageGenerationFunc(height, width, steps, 0, prompt, "", jsonPath, dst, ml, *cfg, appConfig, refImages) + fn, err := backend.ImageGenerationFunc(c.Request().Context(), height, width, steps, 0, prompt, "", jsonPath, dst, ml, *cfg, appConfig, refImages) if err != nil { return err } diff --git a/core/http/endpoints/openai/inpainting_test.go b/core/http/endpoints/openai/inpainting_test.go index 69a80b6de..b251dfda2 100644 --- a/core/http/endpoints/openai/inpainting_test.go +++ b/core/http/endpoints/openai/inpainting_test.go @@ -2,6 +2,7 @@ package openai import ( "bytes" + "context" "mime/multipart" "net/http" "net/http/httptest" @@ -56,7 +57,7 @@ var _ = Describe("Inpainting", func() { appConf := config.NewApplicationConfig(config.WithGeneratedContentDir(tmpDir)) orig := backend.ImageGenerationFunc - backend.ImageGenerationFunc = func(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) { + backend.ImageGenerationFunc = func(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) { fn := func() error { return os.WriteFile(dst, []byte("PNGDATA"), 0644) } diff --git a/core/http/middleware/node_header.go b/core/http/middleware/node_header.go new file mode 100644 index 000000000..96fb4e705 --- /dev/null +++ b/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) + } + } +} diff --git a/core/http/middleware/node_header_concurrency_test.go b/core/http/middleware/node_header_concurrency_test.go new file mode 100644 index 000000000..5f873c9ae --- /dev/null +++ b/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) + } + }) +}) diff --git a/core/http/middleware/node_header_integration_test.go b/core/http/middleware/node_header_integration_test.go new file mode 100644 index 000000000..3f0911752 --- /dev/null +++ b/core/http/middleware/node_header_integration_test.go @@ -0,0 +1,201 @@ +package middleware_test + +// Route-level integration coverage for the X-LocalAI-Node middleware. +// +// What this file pins (and why a separate spec on top of the unit tests +// in node_header_test.go): +// +// - The unit tests in node_header_test.go exercise the wrapper by +// invoking `mw(handler)(c)` directly against a hand-built +// echo.Context. That misses regressions where the contract between +// the real Echo router and the wrapper breaks: e.g. middleware +// installation via e.Use() loses the wrapper because the framework +// re-decorates c.Response().Writer after middleware setup, or a +// handler that bypasses c.Response().Writer (writing to some other +// captured surface). +// +// - This spec dispatches a real HTTP request through e.ServeHTTP into +// a streaming handler shaped like chat.go's streaming branch: set +// SSE headers, write chunks via c.Response().Write, Flush. It +// proves that: +// 1. Middleware installed via e.Use() is on the writer chain +// when the handler runs. +// 2. The per-request holder attached by ExposeNodeHeader is +// visible to the handler (and, transitively, to anything that +// shares the request context, including the SmartRouter). +// 3. The wrapper's lazy maybeSet fires on the first underlying +// Write/Flush, so X-LocalAI-Node lands on the response map +// BEFORE the first body byte is committed. +// 4. The header is present in the recorded response (i.e. it +// isn't dropped because we tried to set it post-WriteHeader). +// +// Out of scope (and why): +// +// - We do NOT wire core/http/endpoints/openai.ChatEndpoint +// end-to-end. ChatEndpoint depends on templates.Evaluator, the +// MCP NATS client, and the LocalAI Assistant holder; standing +// those up just to assert header ordering is out of proportion to +// the property under test. The handler used here mirrors +// chat.go's streaming branch and exercises the SAME middleware -> +// c.Response().Writer -> SSE write path as production. If +// chat.go's streaming branch ever stops going through +// c.Response().Writer (e.g. it starts using a captured raw +// http.ResponseWriter from a different seam), this test will not +// notice; guard that with a code review checklist on chat.go. +// +// - We do NOT spin up the real SmartRouter. The contract between +// router and middleware is "router calls distributedhdr.Stamp on +// the request context; middleware reads the resulting holder on +// first write". A synthetic stamp inside the handler exercises +// the same code path; the router's own unit/integration tests +// cover the routing decision itself. + +import ( + "fmt" + "net/http" + "net/http/httptest" + "strings" + "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/core/http/middleware" + "github.com/mudler/LocalAI/pkg/distributedhdr" +) + +// orderRecorder snapshots the X-LocalAI-Node header value AT THE MOMENT +// the underlying writer is asked to commit each event. Any header set on +// the response map AFTER the first write/flush is dropped on the wire, +// so this is the ground-truth observation a real SSE client would see. +type orderRecorder struct { + http.ResponseWriter + mu sync.Mutex + events []string +} + +func (o *orderRecorder) record(ev string) { + o.mu.Lock() + defer o.mu.Unlock() + o.events = append(o.events, ev) +} + +func (o *orderRecorder) snapshot() []string { + o.mu.Lock() + defer o.mu.Unlock() + out := make([]string, len(o.events)) + copy(out, o.events) + return out +} + +func (o *orderRecorder) WriteHeader(code int) { + o.record(fmt.Sprintf("header:%d:node=%s", code, o.Header().Get(middleware.NodeHeaderName))) + o.ResponseWriter.WriteHeader(code) +} + +func (o *orderRecorder) Write(b []byte) (int, error) { + o.record(fmt.Sprintf("write:node=%s", o.Header().Get(middleware.NodeHeaderName))) + return o.ResponseWriter.Write(b) +} + +func (o *orderRecorder) Flush() { + o.record(fmt.Sprintf("flush:node=%s", o.Header().Get(middleware.NodeHeaderName))) + if f, ok := o.ResponseWriter.(http.Flusher); ok { + f.Flush() + } +} + +var _ = Describe("ExposeNodeHeader middleware (route-level integration)", func() { + const fakeNodeID = "node-route-7" + + var appCfg *config.ApplicationConfig + + BeforeEach(func() { + appCfg = config.NewApplicationConfig() + appCfg.ExposeNodeHeader = true + }) + + It("stamps X-LocalAI-Node before the first SSE byte via the real router + middleware chain", func() { + // Build a real Echo router. We need the tracker to sit BELOW + // the ExposeNodeHeader wrapper in the writer chain (so its + // recorded snapshot reflects what bytes-on-the-wire see AFTER + // the wrapper has had a chance to stamp the header). Install + // the tracker via a middleware that runs BEFORE + // ExposeNodeHeader; Echo's middleware execution order matches + // e.Use() call order, so the first Use() wraps the OUTER + // layer of the writer chain (i.e. the wrapper installed by + // the second Use() wraps the tracker installed by the first). + var ( + recorderMu sync.Mutex + tracker *orderRecorder + ) + e := echo.New() + e.Use(func(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + recorderMu.Lock() + tracker = &orderRecorder{ResponseWriter: c.Response().Writer} + c.Response().Writer = tracker + recorderMu.Unlock() + return next(c) + } + }) + e.Use(middleware.ExposeNodeHeader(appCfg)) + + e.POST("/v1/chat/completions", func(c echo.Context) error { + // Simulate the SmartRouter publishing the picked node ID + // into the per-request holder installed by the middleware. + // In production this happens inside ModelLoader.Load via + // distributedhdr.Stamp(ctx, result.Node.ID). + distributedhdr.Stamp(c.Request().Context(), fakeNodeID) + + // SSE response prelude (same shape as chat.go). + c.Response().Header().Set("Content-Type", "text/event-stream") + c.Response().Header().Set("Cache-Control", "no-cache") + c.Response().Header().Set("Connection", "keep-alive") + + // Emit a handful of SSE chunks. The very first + // Write/Flush is what triggers the middleware + // wrapper's maybeSet, so the X-LocalAI-Node header + // MUST already be on the response map by the time the + // byte is committed. + for i := 0; i < 3; i++ { + _, err := c.Response().Write([]byte(fmt.Sprintf("data: chunk %d\n\n", i))) + if err != nil { + return err + } + c.Response().Flush() + } + return nil + }) + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader("")) + rec := httptest.NewRecorder() + + e.ServeHTTP(rec, req) + + recorderMu.Lock() + Expect(tracker).ToNot(BeNil(), "handler must run and install the order recorder") + events := tracker.snapshot() + recorderMu.Unlock() + + Expect(rec.Code).To(Equal(http.StatusOK)) + Expect(rec.Header().Get(middleware.NodeHeaderName)).To(Equal(fakeNodeID), + "production contract: header must reach the wire on a streamed response") + + Expect(events).ToNot(BeEmpty(), + "expected at least one underlying-writer event from the streaming handler") + + // The very first observed event is the moment the wrapper + // commits to the wire. Its recorded node= value is what a + // real HTTP client would actually see; anything that lands + // AFTER this byte is invisible. + first := events[0] + Expect(first).To(ContainSubstring("node="+fakeNodeID), + "first writer event must carry the X-LocalAI-Node header (chain: middleware.Use -> e.POST -> handler.Write/Flush); got events: %v", events) + + // Body sanity: SSE chunks made it to the recorder. + Expect(rec.Body.String()).To(ContainSubstring("data: chunk 0")) + }) +}) diff --git a/core/http/middleware/node_header_test.go b/core/http/middleware/node_header_test.go new file mode 100644 index 000000000..d7d46adcb --- /dev/null +++ b/core/http/middleware/node_header_test.go @@ -0,0 +1,260 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + + "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" +) + +// orderedWriter records the order in which header-snapshot vs body-byte +// events happen. Used by the streaming spec to assert that the X-LocalAI-Node +// header lands on the response BEFORE the first body byte is committed to +// the underlying writer. +type orderedWriter struct { + http.ResponseWriter + events []string +} + +func (o *orderedWriter) WriteHeader(code int) { + o.events = append(o.events, "header:"+http.StatusText(code)) + o.ResponseWriter.WriteHeader(code) +} + +func (o *orderedWriter) Write(b []byte) (int, error) { + // Snapshot the X-LocalAI-Node header value AT THE INSTANT the underlying + // writer is asked to commit bytes. This is what real HTTP clients + // effectively observe: anything set on the header map AFTER this point + // would be silently dropped. + o.events = append(o.events, "write:node="+o.Header().Get(NodeHeaderName)) + return o.ResponseWriter.Write(b) +} + +func (o *orderedWriter) Flush() { + o.events = append(o.events, "flush:node="+o.Header().Get(NodeHeaderName)) + if f, ok := o.ResponseWriter.(http.Flusher); ok { + f.Flush() + } +} + +var _ = Describe("ExposeNodeHeader middleware", func() { + const ( + fakeNodeID = "node-abcdef" + ) + + var ( + e *echo.Echo + appCfg *config.ApplicationConfig + ) + + BeforeEach(func() { + e = echo.New() + appCfg = &config.ApplicationConfig{} + }) + + // run executes the middleware against a fake handler. The handler may + // reach into the per-request context to stamp the holder (simulating + // what the distributed router does in production); the wrapper reads + // the holder lazily on the first underlying write. + run := func(handler echo.HandlerFunc) *httptest.ResponseRecorder { + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + mw := ExposeNodeHeader(appCfg) + Expect(mw(handler)(c)).To(Succeed()) + return rec + } + + When("ExposeNodeHeader is false", func() { + It("does not set the X-LocalAI-Node header", func() { + appCfg.ExposeNodeHeader = false + + rec := run(func(c echo.Context) error { + // Even if a router were to stamp, with the flag off + // there is no holder on the context so Stamp is a no-op. + distributedhdr.Stamp(c.Request().Context(), fakeNodeID) + return c.String(http.StatusOK, "ok") + }) + + Expect(rec.Header().Get(NodeHeaderName)).To(BeEmpty()) + }) + + It("does not even install the wrapper (writer is unchanged)", func() { + appCfg.ExposeNodeHeader = false + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + origWriter := c.Response().Writer + + handler := func(c echo.Context) error { + // Pass-through must leave the writer identity intact so + // no overhead is added on the hot path when the feature + // is off. + Expect(c.Response().Writer).To(BeIdenticalTo(origWriter)) + // And no holder is attached to the request context. + Expect(distributedhdr.Holder(c.Request().Context())).To(BeNil()) + return c.String(http.StatusOK, "ok") + } + mw := ExposeNodeHeader(appCfg) + Expect(mw(handler)(c)).To(Succeed()) + }) + }) + + When("ExposeNodeHeader is true and the router stamps a node ID", func() { + It("sets the X-LocalAI-Node header on a buffered response", func() { + appCfg.ExposeNodeHeader = true + + rec := run(func(c echo.Context) error { + distributedhdr.Stamp(c.Request().Context(), fakeNodeID) + return c.String(http.StatusOK, "ok") + }) + + Expect(rec.Header().Get(NodeHeaderName)).To(Equal(fakeNodeID)) + }) + + It("sets the header even on a 500 error response (Write still triggers maybeSet)", func() { + appCfg.ExposeNodeHeader = true + + rec := run(func(c echo.Context) error { + distributedhdr.Stamp(c.Request().Context(), fakeNodeID) + return c.String(http.StatusInternalServerError, "boom") + }) + + Expect(rec.Code).To(Equal(http.StatusInternalServerError)) + Expect(rec.Header().Get(NodeHeaderName)).To(Equal(fakeNodeID)) + }) + + It("installs a holder on the request context that the router can find", func() { + appCfg.ExposeNodeHeader = true + + var observed *string + rec := run(func(c echo.Context) error { + h := distributedhdr.Holder(c.Request().Context()) + Expect(h).ToNot(BeNil(), "middleware must attach a per-request holder when the flag is on") + + distributedhdr.Stamp(c.Request().Context(), fakeNodeID) + got := distributedhdr.Load(h) + observed = &got + return c.String(http.StatusOK, "ok") + }) + + Expect(observed).ToNot(BeNil()) + Expect(*observed).To(Equal(fakeNodeID)) + Expect(rec.Header().Get(NodeHeaderName)).To(Equal(fakeNodeID)) + }) + }) + + When("ExposeNodeHeader is true but the router never stamps", func() { + It("does not set the header (in-process model, not distributed)", func() { + appCfg.ExposeNodeHeader = true + + rec := run(func(c echo.Context) error { + // Holder is present but nothing ever stamps it - this is + // the in-process / non-distributed path. + Expect(distributedhdr.Holder(c.Request().Context())).ToNot(BeNil()) + return c.String(http.StatusOK, "ok") + }) + + Expect(rec.Header().Get(NodeHeaderName)).To(BeEmpty()) + }) + }) + + When("the handler streams via Flush before any Write", func() { + It("sets the header BEFORE the first byte hits the underlying writer", func() { + appCfg.ExposeNodeHeader = true + + // Wrap the recorder with an order-tracking writer so we can + // assert that the header is on the response map by the time + // the first body byte is committed. This is the property + // that protected the pre-refactor streaming bug: if the + // wrapper stamped lazily but AFTER the byte commit, real + // SSE clients would see the body without the header. + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + rec := httptest.NewRecorder() + tracker := &orderedWriter{ResponseWriter: rec} + c := e.NewContext(req, rec) + c.Response().Writer = tracker + + handler := func(c echo.Context) error { + // Simulate the router publishing the picked node ID + // mid-request, then an SSE stream emitting chunks. + distributedhdr.Stamp(c.Request().Context(), fakeNodeID) + c.Response().Header().Set("Content-Type", "text/event-stream") + c.Response().Flush() + _, err := c.Response().Write([]byte("data: chunk\n\n")) + return err + } + + mw := ExposeNodeHeader(appCfg) + Expect(mw(handler)(c)).To(Succeed()) + + // First recorded event on the underlying writer must show + // the header already populated. The first event is either + // flush or write; either way the node ID must be on it. + Expect(tracker.events).ToNot(BeEmpty()) + Expect(tracker.events[0]).To(HavePrefix("flush:node=" + fakeNodeID)) + Expect(rec.Header().Get(NodeHeaderName)).To(Equal(fakeNodeID)) + }) + }) + + When("the handler writes a body without an explicit WriteHeader", func() { + It("still stamps the header before the implicit 200 commit", func() { + appCfg.ExposeNodeHeader = true + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + rec := httptest.NewRecorder() + tracker := &orderedWriter{ResponseWriter: rec} + c := e.NewContext(req, rec) + c.Response().Writer = tracker + + handler := func(c echo.Context) error { + distributedhdr.Stamp(c.Request().Context(), fakeNodeID) + _, err := c.Response().Write([]byte("body")) + return err + } + + mw := ExposeNodeHeader(appCfg) + Expect(mw(handler)(c)).To(Succeed()) + + // Echo's Response.Write calls WriteHeader on the underlying + // writer first, then Write. Both must see the header + // already populated (the wrapper's maybeSet ran inside both + // WriteHeader and Write before they hit `tracker`). + Expect(len(tracker.events)).To(BeNumerically(">=", 2)) + Expect(tracker.events[0]).To(HavePrefix("header:")) + Expect(tracker.events[1]).To(Equal("write:node=" + fakeNodeID)) + Expect(rec.Header().Get(NodeHeaderName)).To(Equal(fakeNodeID)) + }) + }) + + When("the router stamps after request entry but before first write", func() { + It("uses the value present AT the first write (late binding)", func() { + appCfg.ExposeNodeHeader = true + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + + handler := func(c echo.Context) error { + // Simulate the router making a routing decision after + // the handler has already started running but before the + // first byte hits the wire. The wrapper must read the + // holder lazily, not eagerly at request entry. + distributedhdr.Stamp(c.Request().Context(), "fresh-node-B") + return c.String(http.StatusOK, "ok") + } + + mw := ExposeNodeHeader(appCfg) + Expect(mw(handler)(c)).To(Succeed()) + + Expect(rec.Header().Get(NodeHeaderName)).To(Equal("fresh-node-B"), + "the wrapper must read the node ID lazily at first write, not eagerly at entry") + }) + }) + +}) diff --git a/core/http/middleware/request.go b/core/http/middleware/request.go index 456c4d0f7..7f3df885e 100644 --- a/core/http/middleware/request.go +++ b/core/http/middleware/request.go @@ -14,6 +14,7 @@ import ( "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/templates" + "github.com/mudler/LocalAI/pkg/distributedhdr" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/utils" "github.com/mudler/xlog" @@ -220,6 +221,7 @@ func (re *RequestExtractor) SetOpenAIRequest(c echo.Context) error { // Add the correlation ID to the new context ctxWithCorrelationID := context.WithValue(c1, CorrelationIDKey, correlationID) + ctxWithCorrelationID = distributedhdr.Inherit(ctxWithCorrelationID, reqCtx) input.Context = ctxWithCorrelationID input.Cancel = cancel @@ -632,6 +634,7 @@ func (re *RequestExtractor) SetOpenResponsesRequest(c echo.Context) error { // Add the correlation ID to the new context ctxWithCorrelationID := context.WithValue(c1, CorrelationIDKey, correlationID) + ctxWithCorrelationID = distributedhdr.Inherit(ctxWithCorrelationID, reqCtx) input.Context = ctxWithCorrelationID input.Cancel = cancel diff --git a/core/http/routes/anthropic.go b/core/http/routes/anthropic.go index 386942e27..e68f88d53 100644 --- a/core/http/routes/anthropic.go +++ b/core/http/routes/anthropic.go @@ -16,6 +16,7 @@ import ( "github.com/mudler/LocalAI/core/services/routing/pii" "github.com/mudler/LocalAI/core/services/routing/piiadapter" "github.com/mudler/LocalAI/core/services/routing/router" + "github.com/mudler/LocalAI/pkg/distributedhdr" "github.com/mudler/xlog" ) @@ -40,6 +41,7 @@ func RegisterAnthropicRoutes(app *echo.Echo, ) messagesMiddleware := []echo.MiddlewareFunc{ + middleware.ExposeNodeHeader(application.ApplicationConfig()), middleware.UsageMiddleware(application.StatsRecorder(), application.FallbackUser()), middleware.TraceMiddleware(application), re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_CHAT)), @@ -111,6 +113,7 @@ func setAnthropicRequestContext(appConfig *config.ApplicationConfig) echo.Middle // Add the correlation ID to the new context ctxWithCorrelationID := context.WithValue(c1, middleware.CorrelationIDKey, correlationID) + ctxWithCorrelationID = distributedhdr.Inherit(ctxWithCorrelationID, reqCtx) input.Context = ctxWithCorrelationID input.Cancel = cancel diff --git a/core/http/routes/jina.go b/core/http/routes/jina.go index b4fafbc57..dbc377412 100644 --- a/core/http/routes/jina.go +++ b/core/http/routes/jina.go @@ -20,6 +20,7 @@ func RegisterJINARoutes(app *echo.Echo, rerankHandler := jina.JINARerankEndpoint(cl, ml, appConfig) app.POST("/v1/rerank", rerankHandler, + middleware.ExposeNodeHeader(appConfig), re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_RERANK)), re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.JINARerankRequest) })) } diff --git a/core/http/routes/localai.go b/core/http/routes/localai.go index 4f61c8525..96baceaf8 100644 --- a/core/http/routes/localai.go +++ b/core/http/routes/localai.go @@ -165,12 +165,15 @@ func RegisterLocalAIRoutes(router *echo.Echo, middleware.TraceMiddleware(app)) vadHandler := localai.VADEndpoint(cl, ml, appConfig) + vadNodeHeader := middleware.ExposeNodeHeader(appConfig) router.POST("/vad", vadHandler, + vadNodeHeader, requestExtractor.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_VAD)), requestExtractor.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.VADRequest) })) router.POST("/v1/vad", vadHandler, + vadNodeHeader, requestExtractor.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_VAD)), requestExtractor.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.VADRequest) })) diff --git a/core/http/routes/ollama.go b/core/http/routes/ollama.go index f02db7636..4145dd4f8 100644 --- a/core/http/routes/ollama.go +++ b/core/http/routes/ollama.go @@ -10,6 +10,7 @@ import ( "github.com/mudler/LocalAI/core/http/endpoints/ollama" "github.com/mudler/LocalAI/core/http/middleware" "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/pkg/distributedhdr" ) func RegisterOllamaRoutes(app *echo.Echo, @@ -18,6 +19,7 @@ func RegisterOllamaRoutes(app *echo.Echo, traceMiddleware := middleware.TraceMiddleware(application) usageMiddleware := middleware.UsageMiddleware(application.StatsRecorder(), application.FallbackUser()) + nodeHeaderMiddleware := middleware.ExposeNodeHeader(application.ApplicationConfig()) // Chat endpoint: POST /api/chat chatHandler := ollama.ChatEndpoint( @@ -27,6 +29,7 @@ func RegisterOllamaRoutes(app *echo.Echo, application.ApplicationConfig(), ) chatMiddleware := []echo.MiddlewareFunc{ + nodeHeaderMiddleware, usageMiddleware, traceMiddleware, re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_CHAT)), @@ -43,6 +46,7 @@ func RegisterOllamaRoutes(app *echo.Echo, application.ApplicationConfig(), ) generateMiddleware := []echo.MiddlewareFunc{ + nodeHeaderMiddleware, usageMiddleware, traceMiddleware, re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_CHAT)), @@ -58,6 +62,7 @@ func RegisterOllamaRoutes(app *echo.Echo, application.ApplicationConfig(), ) embedMiddleware := []echo.MiddlewareFunc{ + nodeHeaderMiddleware, usageMiddleware, traceMiddleware, re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_EMBEDDINGS)), @@ -108,6 +113,7 @@ func setOllamaChatRequestContext(appConfig *config.ApplicationConfig) echo.Middl }() ctxWithCorrelationID := context.WithValue(c1, middleware.CorrelationIDKey, correlationID) + ctxWithCorrelationID = distributedhdr.Inherit(ctxWithCorrelationID, reqCtx) input.Context = ctxWithCorrelationID input.Cancel = cancel @@ -149,6 +155,7 @@ func setOllamaGenerateRequestContext(appConfig *config.ApplicationConfig) echo.M }() ctxWithCorrelationID := context.WithValue(c1, middleware.CorrelationIDKey, correlationID) + ctxWithCorrelationID = distributedhdr.Inherit(ctxWithCorrelationID, reqCtx) input.Ctx = ctxWithCorrelationID input.Cancel = cancel diff --git a/core/http/routes/openai.go b/core/http/routes/openai.go index 417df8a79..8a13935ae 100644 --- a/core/http/routes/openai.go +++ b/core/http/routes/openai.go @@ -20,6 +20,13 @@ func RegisterOpenAIRoutes(app *echo.Echo, // openAI compatible API endpoint traceMiddleware := middleware.TraceMiddleware(application) usageMiddleware := middleware.UsageMiddleware(application.StatsRecorder(), application.FallbackUser()) + // X-LocalAI-Node attribution middleware: wraps the response writer and + // stamps the header on first write when --expose-node-header is on. No-op + // otherwise. Applied to every inference path that routes through + // ml.Load (chat, completion, embeddings, audio transcriptions/speech, + // image generation/inpainting) so distributed-mode operators can observe + // which worker served each request. + nodeHeaderMiddleware := middleware.ExposeNodeHeader(application.ApplicationConfig()) // realtime // TODO: Modify/disable the API key middleware for this endpoint to allow ephemeral keys created by sessions @@ -37,6 +44,7 @@ func RegisterOpenAIRoutes(app *echo.Echo, // chat chatHandler := openai.ChatEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.TemplatesEvaluator(), application.ApplicationConfig(), natsClient, application.LocalAIAssistant(), application.PIIRedactor(), application.PIIEvents()) chatMiddleware := []echo.MiddlewareFunc{ + nodeHeaderMiddleware, usageMiddleware, traceMiddleware, re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_CHAT)), @@ -110,6 +118,7 @@ func RegisterOpenAIRoutes(app *echo.Echo, // completion completionHandler := openai.CompletionEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.TemplatesEvaluator(), application.ApplicationConfig(), application.PIIRedactor(), application.PIIEvents()) completionMiddleware := []echo.MiddlewareFunc{ + nodeHeaderMiddleware, usageMiddleware, traceMiddleware, re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_COMPLETION)), @@ -131,6 +140,7 @@ func RegisterOpenAIRoutes(app *echo.Echo, // embeddings embeddingHandler := openai.EmbeddingsEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig()) embeddingMiddleware := []echo.MiddlewareFunc{ + nodeHeaderMiddleware, usageMiddleware, traceMiddleware, re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_EMBEDDINGS)), @@ -151,6 +161,7 @@ func RegisterOpenAIRoutes(app *echo.Echo, audioHandler := openai.TranscriptEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig()) audioMiddleware := []echo.MiddlewareFunc{ + nodeHeaderMiddleware, traceMiddleware, re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_TRANSCRIPT)), re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }), @@ -186,6 +197,7 @@ func RegisterOpenAIRoutes(app *echo.Echo, audioSpeechHandler := localai.TTSEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig()) audioSpeechMiddleware := []echo.MiddlewareFunc{ + nodeHeaderMiddleware, traceMiddleware, re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_TTS)), re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.TTSRequest) }), @@ -197,6 +209,7 @@ func RegisterOpenAIRoutes(app *echo.Echo, // images imageHandler := openai.ImageEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig()) imageMiddleware := []echo.MiddlewareFunc{ + nodeHeaderMiddleware, traceMiddleware, // Default: use the first available image generation model re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_IMAGE)), diff --git a/core/services/nodes/model_router.go b/core/services/nodes/model_router.go index d834bcbc8..2f87bc5e6 100644 --- a/core/services/nodes/model_router.go +++ b/core/services/nodes/model_router.go @@ -5,6 +5,7 @@ import ( "fmt" "sync" + "github.com/mudler/LocalAI/pkg/distributedhdr" pb "github.com/mudler/LocalAI/pkg/grpc/proto" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/xlog" @@ -69,6 +70,16 @@ func (a *ModelRouterAdapter) Route(ctx context.Context, backend, modelID, modelN // the ModelLoader returns this model on subsequent requests. m := model.NewModelWithClient(modelID, result.Node.Address, result.Client) + // Publish the picked node ID into the per-request holder attached to + // ctx (by middleware.ExposeNodeHeader). No-op when the holder is + // absent - e.g. health-check probes, watchdog reloads or any path that + // did not flow through the HTTP middleware. The response writer + // wrapper reads this on the first byte to set X-LocalAI-Node, so + // attribution is exact for the request that triggered THIS Route call, + // even when several requests for the same modelID are being routed + // concurrently to different replicas. + distributedhdr.Stamp(ctx, result.Node.ID) + xlog.Info("Model routed to remote node", "model", modelName, "node", result.Node.Name, "address", result.Node.Address) return m, nil } diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 792539bce..03a251b78 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -88,6 +88,7 @@ The frontend is a standard LocalAI instance with distributed mode enabled. These | `--auth-database-url` | `LOCALAI_AUTH_DATABASE_URL` | *(required)* | PostgreSQL connection URL | | `--backend-install-timeout` | `LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT` | `15m` | How long the frontend waits for a worker to acknowledge a backend install before considering the request stalled. Raise it when workers pull large backend images over slow links. If a worker takes longer than this, the operation shows as "still installing in background" in the admin UI and clears once the worker finishes. | | `--backend-upgrade-timeout` | `LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT` | `15m` | Same as the install timeout, applied to backend upgrades (force-reinstall). | +| `--expose-node-header` | `LOCALAI_EXPOSE_NODE_HEADER` | `false` | When enabled, inference responses carry an `X-LocalAI-Node` header with the ID of the worker node that served the request. Coverage spans the OpenAI-compatible endpoints (chat completions, completions, embeddings, audio transcriptions, audio speech / TTS, image generations, image inpainting), the Jina rerank endpoint (`/v1/rerank`), the VAD endpoints (`/v1/vad`, `/vad`), and the Anthropic Messages (`/v1/messages`) and Ollama (`/api/chat`, `/api/generate`, `/api/embed`) shims. Useful for debugging, observability and load-balancer attribution. Off by default: the node ID reveals internal cluster topology and should not be exposed on a public endpoint. Best-effort: under heavy concurrency for the same model across multiple replicas, the header may reflect a recent routing decision rather than this exact request's. Acceptable for observability and debugging. | ### Optional: S3 Object Storage diff --git a/pkg/distributedhdr/distributedhdr.go b/pkg/distributedhdr/distributedhdr.go new file mode 100644 index 000000000..4f889c633 --- /dev/null +++ b/pkg/distributedhdr/distributedhdr.go @@ -0,0 +1,108 @@ +// Package distributedhdr carries a per-request "which worker node served +// me" record from the distributed router (core/services/nodes) up to the +// HTTP response writer wrapper (core/http/middleware). It is the bridge +// for the X-LocalAI-Node response header without coupling those two +// packages directly or going through any shared mutable state. +// +// Why its own package: both core/http/middleware and core/services/nodes +// already import pkg/model, and the natural homes for this key would +// create an import cycle if either side hosted the helper. Putting the +// key and the tiny helpers here keeps it neutral - producer and consumer +// import a leaf package, not each other. +package distributedhdr + +import ( + "context" + "sync/atomic" +) + +// ctxKey is an unexported context-key type so external packages cannot +// collide with our key by accident. +type ctxKey struct{} + +// holderKey is the singleton context key used by WithHolder / Holder / +// Stamp. Unexported on purpose: producers and consumers must go through +// the helpers below so the storage representation stays an +// implementation detail. +var holderKey = ctxKey{} + +// NewHolder returns an empty holder ready to be attached to a request +// context via WithHolder. The middleware creates one per HTTP request; +// the router fills it; the response writer wrapper reads it on the +// first byte. The atomic.Value gives us race-clean publish across the +// goroutines that may write (router) and read (response writer +// wrapper) the slot. +func NewHolder() *atomic.Value { + return &atomic.Value{} +} + +// WithHolder returns a derived context that carries the provided holder. +// The middleware calls this on the per-request context BEFORE the +// downstream handler chain runs, so any goroutine that inherits this +// context (notably the SmartRouter / ModelRouterAdapter) can find the +// holder via Stamp. +func WithHolder(ctx context.Context, h *atomic.Value) context.Context { + if h == nil { + return ctx + } + return context.WithValue(ctx, holderKey, h) +} + +// Holder returns the holder attached to ctx, or nil if none was set. +// Callers that just want to publish should prefer Stamp; Holder is +// useful for tests and for propagating the holder across derived +// contexts (see Inherit). +func Holder(ctx context.Context) *atomic.Value { + if ctx == nil { + return nil + } + h, _ := ctx.Value(holderKey).(*atomic.Value) + return h +} + +// Stamp records the picked worker node ID into the holder attached to +// ctx. No-op when: +// - ctx is nil +// - no holder is attached (e.g. the X-LocalAI-Node feature is off, so +// the middleware never ran) +// - nodeID is empty +// +// Stamp is safe to call from any goroutine. Subsequent reads via Load +// observe the most recent stamp. +func Stamp(ctx context.Context, nodeID string) { + if nodeID == "" { + return + } + h := Holder(ctx) + if h == nil { + return + } + h.Store(nodeID) +} + +// Load returns the node ID most recently stamped into the holder, or +// "" if nothing has been stamped. Intended for the response writer +// wrapper's first-byte hook. +func Load(h *atomic.Value) string { + if h == nil { + return "" + } + v, _ := h.Load().(string) + return v +} + +// Inherit copies the holder reference from src into dst when present. +// Used at request-handling seams where the request context is replaced +// with a fresh context derived from the long-lived application context +// (so the cancel semantics of the original context are preserved while +// also allowing the load path to outlive the HTTP transport). The +// holder is a pointer, so both contexts point at the same slot; a +// router stamp via Stamp(dst, ...) is observed by the middleware that +// reads through src's holder. +func Inherit(dst, src context.Context) context.Context { + h := Holder(src) + if h == nil { + return dst + } + return WithHolder(dst, h) +} diff --git a/pkg/distributedhdr/distributedhdr_test.go b/pkg/distributedhdr/distributedhdr_test.go new file mode 100644 index 000000000..3a90938f1 --- /dev/null +++ b/pkg/distributedhdr/distributedhdr_test.go @@ -0,0 +1,99 @@ +package distributedhdr_test + +import ( + "context" + "sync" + "testing" + + "github.com/mudler/LocalAI/pkg/distributedhdr" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestDistributedHdr(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "distributedhdr suite") +} + +var _ = Describe("distributedhdr", func() { + It("returns empty string when no holder is attached", func() { + Expect(distributedhdr.Holder(context.Background())).To(BeNil()) + Expect(distributedhdr.Load(distributedhdr.Holder(context.Background()))).To(BeEmpty()) + }) + + It("is a no-op when stamping into a context without a holder", func() { + Expect(func() { + distributedhdr.Stamp(context.Background(), "worker-x") + }).ToNot(Panic()) + }) + + It("is a no-op when stamping a nil context", func() { + Expect(func() { + distributedhdr.Stamp(nil, "worker-x") + }).ToNot(Panic()) + }) + + It("stores and retrieves a node ID through Stamp/Load", func() { + h := distributedhdr.NewHolder() + ctx := distributedhdr.WithHolder(context.Background(), h) + + distributedhdr.Stamp(ctx, "worker-a") + Expect(distributedhdr.Load(h)).To(Equal("worker-a")) + }) + + It("retains the latest stamp when called twice", func() { + h := distributedhdr.NewHolder() + ctx := distributedhdr.WithHolder(context.Background(), h) + + distributedhdr.Stamp(ctx, "worker-a") + distributedhdr.Stamp(ctx, "worker-b") + Expect(distributedhdr.Load(h)).To(Equal("worker-b")) + }) + + It("does not overwrite when stamping with an empty string", func() { + h := distributedhdr.NewHolder() + ctx := distributedhdr.WithHolder(context.Background(), h) + + distributedhdr.Stamp(ctx, "worker-a") + distributedhdr.Stamp(ctx, "") + Expect(distributedhdr.Load(h)).To(Equal("worker-a")) + }) + + It("propagates the holder via Inherit into a derived context", func() { + h := distributedhdr.NewHolder() + src := distributedhdr.WithHolder(context.Background(), h) + dst := distributedhdr.Inherit(context.Background(), src) + + distributedhdr.Stamp(dst, "worker-c") + Expect(distributedhdr.Load(h)).To(Equal("worker-c")) + }) + + It("Inherit is a no-op when src has no holder", func() { + dst := distributedhdr.Inherit(context.Background(), context.Background()) + Expect(distributedhdr.Holder(dst)).To(BeNil()) + }) + + It("publishes stamps across goroutines race-free", func() { + h := distributedhdr.NewHolder() + ctx := distributedhdr.WithHolder(context.Background(), h) + + const N = 64 + var wg sync.WaitGroup + wg.Add(N) + for i := 0; i < N; i++ { + i := i + go func() { + defer wg.Done() + if i%2 == 0 { + distributedhdr.Stamp(ctx, "worker-even") + } else { + distributedhdr.Stamp(ctx, "worker-odd") + } + }() + } + wg.Wait() + + got := distributedhdr.Load(h) + Expect(got).To(SatisfyAny(Equal("worker-even"), Equal("worker-odd"))) + }) +})