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"))) + }) +})