|
| 1 | +package backend_test |
| 2 | + |
| 3 | +// Regression spec for X-LocalAI-Node coverage on audio/image/TTS/rerank/VAD. |
| 4 | +// |
| 5 | +// The X-LocalAI-Node middleware (core/http/middleware.ExposeNodeHeader) |
| 6 | +// works end-to-end only if the per-request holder attached to the HTTP |
| 7 | +// request context reaches the SmartRouter via ml.Load(opts...). The chain |
| 8 | +// is: |
| 9 | +// |
| 10 | +// handler -> backend.Foo(ctx, ...) -> ModelOptions(cfg, app, WithContext(ctx)) |
| 11 | +// -> ml.Load(opts...) -> grpcModel(..., o.context) -> modelRouter(ctx, ...) |
| 12 | +// -> SmartRouter -> distributedhdr.Stamp(ctx, nodeID) |
| 13 | +// |
| 14 | +// If any backend helper drops `ctx` and lets ModelOptions fall back to the |
| 15 | +// app context, the router never sees the per-request holder and the |
| 16 | +// header silently stays empty for that endpoint. These specs pin the |
| 17 | +// request-context-reaches-router contract for the five backend helpers |
| 18 | +// that were previously dropping ctx between the handler and Load. |
| 19 | + |
| 20 | +import ( |
| 21 | + "context" |
| 22 | + "sync/atomic" |
| 23 | + |
| 24 | + "github.com/mudler/LocalAI/core/backend" |
| 25 | + "github.com/mudler/LocalAI/core/config" |
| 26 | + "github.com/mudler/LocalAI/core/schema" |
| 27 | + pbproto "github.com/mudler/LocalAI/pkg/grpc/proto" |
| 28 | + "github.com/mudler/LocalAI/pkg/distributedhdr" |
| 29 | + "github.com/mudler/LocalAI/pkg/model" |
| 30 | + "github.com/mudler/LocalAI/pkg/system" |
| 31 | + |
| 32 | + . "github.com/onsi/ginkgo/v2" |
| 33 | + . "github.com/onsi/gomega" |
| 34 | +) |
| 35 | + |
| 36 | +// newCapturingLoader returns a ModelLoader wired with a stub model router |
| 37 | +// that captures the context it receives and then short-circuits with a |
| 38 | +// sentinel error. The router callback is the exact seam where the |
| 39 | +// SmartRouter would call distributedhdr.Stamp in production, so observing |
| 40 | +// the holder here is equivalent to observing it at the real router. |
| 41 | +func newCapturingLoader() (*model.ModelLoader, *atomic.Value, func() context.Context) { |
| 42 | + loader := model.NewModelLoader(&system.SystemState{}) |
| 43 | + var captured atomic.Value |
| 44 | + loader.SetModelRouter(func(ctx context.Context, _ string, _, _, _ string, _ *pbproto.ModelOptions, _ bool) (*model.Model, error) { |
| 45 | + captured.Store(ctx) |
| 46 | + // Return an error so the backend short-circuits before trying to |
| 47 | + // dial gRPC. We only care about the context-arrival contract. |
| 48 | + return nil, errRouterShortCircuit |
| 49 | + }) |
| 50 | + get := func() context.Context { |
| 51 | + v, _ := captured.Load().(context.Context) |
| 52 | + return v |
| 53 | + } |
| 54 | + return loader, &captured, get |
| 55 | +} |
| 56 | + |
| 57 | +var errRouterShortCircuit = sentinelErr("router short-circuit (test)") |
| 58 | + |
| 59 | +type sentinelErr string |
| 60 | + |
| 61 | +func (s sentinelErr) Error() string { return string(s) } |
| 62 | + |
| 63 | +func newAppCfg() *config.ApplicationConfig { |
| 64 | + return config.NewApplicationConfig(config.WithSystemState(&system.SystemState{})) |
| 65 | +} |
| 66 | + |
| 67 | +func newModelCfg() config.ModelConfig { |
| 68 | + threads := 1 |
| 69 | + cfg := config.ModelConfig{ |
| 70 | + Name: "test-model", |
| 71 | + Backend: "stub-backend", |
| 72 | + Threads: &threads, |
| 73 | + } |
| 74 | + cfg.Model = "test.bin" |
| 75 | + return cfg |
| 76 | +} |
| 77 | + |
| 78 | +var _ = Describe("X-LocalAI-Node ctx propagation contract", func() { |
| 79 | + const fakeNodeID = "node-ctx-propagation-7" |
| 80 | + |
| 81 | + var ( |
| 82 | + appCfg *config.ApplicationConfig |
| 83 | + modelCfg config.ModelConfig |
| 84 | + loader *model.ModelLoader |
| 85 | + routerCtxOf func() context.Context |
| 86 | + holder *atomic.Value |
| 87 | + reqCtx context.Context |
| 88 | + ) |
| 89 | + |
| 90 | + BeforeEach(func() { |
| 91 | + appCfg = newAppCfg() |
| 92 | + modelCfg = newModelCfg() |
| 93 | + loader, _, routerCtxOf = newCapturingLoader() |
| 94 | + holder = distributedhdr.NewHolder() |
| 95 | + reqCtx = distributedhdr.WithHolder(context.Background(), holder) |
| 96 | + }) |
| 97 | + |
| 98 | + // stampViaRouterCtx asserts the captured router context carries the |
| 99 | + // SAME holder that was attached to the request. We verify by stamping |
| 100 | + // through the router-side ctx and observing the value via the |
| 101 | + // request-side holder; if the holders were different objects the load |
| 102 | + // would return "". |
| 103 | + stampViaRouterCtx := func() { |
| 104 | + routerCtx := routerCtxOf() |
| 105 | + Expect(routerCtx).ToNot(BeNil(), "router callback must have been invoked") |
| 106 | + distributedhdr.Stamp(routerCtx, fakeNodeID) |
| 107 | + Expect(distributedhdr.Load(holder)).To(Equal(fakeNodeID), |
| 108 | + "stamp via router-side ctx must be observable via the request-side holder") |
| 109 | + } |
| 110 | + |
| 111 | + It("Rerank forwards the request context to the SmartRouter", func() { |
| 112 | + _, err := backend.Rerank(reqCtx, &pbproto.RerankRequest{Query: "q"}, loader, appCfg, modelCfg) |
| 113 | + Expect(err).To(HaveOccurred()) |
| 114 | + Expect(err.Error()).To(ContainSubstring("router short-circuit (test)")) |
| 115 | + stampViaRouterCtx() |
| 116 | + }) |
| 117 | + |
| 118 | + It("VAD forwards the request context to the SmartRouter", func() { |
| 119 | + _, err := backend.VAD(&schema.VADRequest{}, reqCtx, loader, appCfg, modelCfg) |
| 120 | + Expect(err).To(HaveOccurred()) |
| 121 | + Expect(err.Error()).To(ContainSubstring("router short-circuit (test)")) |
| 122 | + stampViaRouterCtx() |
| 123 | + }) |
| 124 | + |
| 125 | + It("ModelTTS forwards the request context to the SmartRouter", func() { |
| 126 | + _, _, err := backend.ModelTTS(reqCtx, "hello", "", "", loader, appCfg, modelCfg) |
| 127 | + Expect(err).To(HaveOccurred()) |
| 128 | + Expect(err.Error()).To(ContainSubstring("router short-circuit (test)")) |
| 129 | + stampViaRouterCtx() |
| 130 | + }) |
| 131 | + |
| 132 | + It("ModelTTSStream forwards the request context to the SmartRouter", func() { |
| 133 | + err := backend.ModelTTSStream(reqCtx, "hello", "", "", loader, appCfg, modelCfg, func([]byte) error { return nil }) |
| 134 | + Expect(err).To(HaveOccurred()) |
| 135 | + Expect(err.Error()).To(ContainSubstring("router short-circuit (test)")) |
| 136 | + stampViaRouterCtx() |
| 137 | + }) |
| 138 | + |
| 139 | + It("ModelTranscriptionWithOptions forwards the request context to the SmartRouter", func() { |
| 140 | + _, err := backend.ModelTranscriptionWithOptions(reqCtx, backend.TranscriptionRequest{Audio: "x.wav"}, loader, modelCfg, appCfg) |
| 141 | + Expect(err).To(HaveOccurred()) |
| 142 | + Expect(err.Error()).To(ContainSubstring("router short-circuit (test)")) |
| 143 | + stampViaRouterCtx() |
| 144 | + }) |
| 145 | + |
| 146 | + It("ModelTranscriptionStream forwards the request context to the SmartRouter", func() { |
| 147 | + err := backend.ModelTranscriptionStream(reqCtx, backend.TranscriptionRequest{Audio: "x.wav"}, loader, modelCfg, appCfg, func(backend.TranscriptionStreamChunk) {}) |
| 148 | + Expect(err).To(HaveOccurred()) |
| 149 | + Expect(err.Error()).To(ContainSubstring("router short-circuit (test)")) |
| 150 | + stampViaRouterCtx() |
| 151 | + }) |
| 152 | + |
| 153 | + It("ImageGeneration forwards the request context to the SmartRouter", func() { |
| 154 | + _, err := backend.ImageGeneration(reqCtx, 64, 64, 1, 0, "p", "", "", "/tmp/out.png", loader, modelCfg, appCfg, nil) |
| 155 | + Expect(err).To(HaveOccurred()) |
| 156 | + Expect(err.Error()).To(ContainSubstring("router short-circuit (test)")) |
| 157 | + stampViaRouterCtx() |
| 158 | + }) |
| 159 | + |
| 160 | + It("does NOT leak the holder when the app context is used instead", func() { |
| 161 | + // Sanity: the bug being fixed manifests as the router getting |
| 162 | + // appCfg.Context (no holder) instead of reqCtx (holder). A direct |
| 163 | + // call with context.Background() must not see the holder via the |
| 164 | + // app context surface. |
| 165 | + appCtxOnly := appCfg.Context |
| 166 | + Expect(distributedhdr.Holder(appCtxOnly)).To(BeNil(), |
| 167 | + "the app context must not be the carrier of per-request holders") |
| 168 | + }) |
| 169 | +}) |
0 commit comments