LocalAI/core/http/middleware/diagnostics_test.go

// SPDX-License-Identifier: MIT
package middleware_test

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"

	"github.com/google/uuid"
	"github.com/labstack/echo/v4"
	"github.com/mudler/LocalAI/core/config"
	. "github.com/mudler/LocalAI/core/http/middleware"
	"github.com/mudler/LocalAI/core/schema"
	"github.com/mudler/LocalAI/core/services/failover"
	"github.com/mudler/LocalAI/pkg/diagnostics"
	"github.com/mudler/LocalAI/pkg/model"
	"github.com/mudler/LocalAI/pkg/system"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
)

var _ = Describe("Diagnostics extraction", func() {
	var re *RequestExtractor
	var ac *config.ApplicationConfig
	var bcl *config.ModelConfigLoader
	var events []diagnostics.Event
	BeforeEach(func() {
		dir := GinkgoT().TempDir()
		ac = config.NewApplicationConfig()
		ac.Context = context.Background()
		ac.SystemState = &system.SystemState{Model: system.Model{ModelsPath: dir}}
		bcl = config.NewModelConfigLoader(dir)
		a := config.ModelConfig{Name: "sentinel-model"}
		a.Model = "sentinel-model"
		bcl.ReplaceModelConfigs([]config.ModelConfig{a})
		events = nil
		ac.DiagnosticsRecorder = diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) })
		re = NewRequestExtractor(bcl, model.NewModelLoader(ac.SystemState), ac)
	})
	ends := func(p diagnostics.Phase) []diagnostics.Event {
		var out []diagnostics.Event
		for _, e := range events {
			if e.Phase == p && e.State == diagnostics.StateEnd {
				out = append(out, e)
			}
		}
		return out
	}
	It("skips chat fallback discovery for a bound model but retains body validation", func() {
		e := echo.New()
		e.POST("/", func(c echo.Context) error { return c.NoContent(204) },
			re.SetModelAndConfigWithDefault(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }, nil))
		req := httptest.NewRequest("POST", "/", strings.NewReader(`{"model":"sentinel-model"}`))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Authorization", "Bearer not-a-model")
		rec := httptest.NewRecorder()
		e.ServeHTTP(rec, req)
		Expect(rec.Code).To(Equal(204))
		Expect(ends(diagnostics.PhaseBearerLookup)).To(BeEmpty())
		Expect(ends(diagnostics.PhaseDefaultListing)).To(BeEmpty())
		Expect(ends(diagnostics.PhaseBodyLookup)).To(HaveLen(1))
	})
	It("preserves lazy fallback precedence and malformed JSON behavior", func() {
		for _, tc := range []struct {
			body, path, query, bearer, local string
			listing, lookup, code            int
		}{
			{body: `{}`, listing: 1, code: 204},
			{body: `{}`, bearer: "sentinel-model", lookup: 1, code: 204},
			{body: `{}`, bearer: "not-a-model", lookup: 1, listing: 1, code: 204},
			{body: `{}`, query: "?model=sentinel-model", bearer: "not-a-model", code: 204},
			{body: `{}`, path: "sentinel-model", query: "?model=missing", code: 204},
			{body: `{}`, local: "sentinel-model", path: "missing", code: 204},
			{body: `{"model":"sentinel-model"}`, bearer: "missing", local: "missing", code: 204},
			{body: `{`, bearer: "sentinel-model", code: 400},
			{body: `{"model":"missing"}`, bearer: "sentinel-model", code: 404},
		} {
			events = nil
			e := echo.New()
			r := httptest.NewRequest("POST", "/"+tc.query, strings.NewReader(tc.body))
			r.Header.Set("Content-Type", "application/json")
			if tc.bearer != "" {
				r.Header.Set("Authorization", "Bearer "+tc.bearer)
			}
			w := httptest.NewRecorder()
			c := e.NewContext(r, w)
			c.SetParamNames("model")
			c.SetParamValues(tc.path)
			c.Set(CONTEXT_LOCALS_KEY_MODEL_NAME, tc.local)
			h := re.SetModelAndConfigWithDefault(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }, nil)(func(c echo.Context) error {
				Expect(c.Get(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenAIRequest).Model).To(Equal("sentinel-model"))
				return c.NoContent(204)
			})
			if err := h(c); err != nil {
				e.HTTPErrorHandler(err, c)
			}
			Expect(w.Code).To(Equal(tc.code))
			Expect(ends(diagnostics.PhaseDefaultListing)).To(HaveLen(tc.listing))
			Expect(ends(diagnostics.PhaseBearerLookup)).To(HaveLen(tc.lookup))
		}
	})
	It("retains alias, disabled, regex existence, and form behavior on the lazy path", func() {
		disabled := true
		a := config.ModelConfig{Name: "alias", Alias: "sentinel-model"}
		d := config.ModelConfig{Name: "disabled", Disabled: &disabled}
		d.Model = d.Name
		target := config.ModelConfig{Name: "sentinel-model"}
		target.Model = target.Name
		bcl.ReplaceModelConfigs([]config.ModelConfig{target, a, d})
		// Run the unchanged eager API as the compatibility reference as well.
		for _, lazy := range []bool{false, true} {
			for _, tc := range []struct {
				body, contentType, served string
				code                      int
			}{
				{`{"model":"alias"}`, "application/json", "sentinel-model", 204},
				{`{"model":"disabled"}`, "application/json", "", 403},
				{`{"model":"sentinel-.*"}`, "application/json", "", 204},
				{"model=sentinel-model", "application/x-www-form-urlencoded", "sentinel-model", 204},
			} {
				events = nil
				e := echo.New()
				chain := []echo.MiddlewareFunc{re.BuildFilteredFirstAvailableDefaultModel(nil), re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) })}
				if lazy {
					chain = []echo.MiddlewareFunc{re.SetModelAndConfigWithDefault(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }, nil)}
				}
				e.POST("/", func(c echo.Context) error {
					Expect(c.Get(CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig).Name).To(Equal(tc.served))
					return c.NoContent(204)
				}, chain...)
				r := httptest.NewRequest("POST", "/", strings.NewReader(tc.body))
				r.Header.Set("Content-Type", tc.contentType)
				r.Header.Set("Authorization", "Bearer missing")
				w := httptest.NewRecorder()
				e.ServeHTTP(w, r)
				Expect(w.Code).To(Equal(tc.code))
				if lazy {
					Expect(ends(diagnostics.PhaseBearerLookup)).To(BeEmpty())
					Expect(ends(diagnostics.PhaseDefaultListing)).To(BeEmpty())
				}
			}
		}
	})

	It("sets the same ID before a streaming flush and emits no header when disabled", func() {
		recorder := ac.DiagnosticsRecorder
		for _, enabled := range []bool{false, true} {
			ac.DiagnosticsRecorder = nil
			if enabled {
				ac.DiagnosticsRecorder = recorder
			}
			e := echo.New()
			e.POST("/", func(c echo.Context) error {
				id := diagnostics.RequestID(c.Request().Context())
				Expect(c.Response().Header().Get(diagnostics.HeaderDiagnosticID)).To(Equal(id))
				Expect(id != "").To(Equal(enabled))
				c.Response().WriteHeader(200)
				c.Response().Flush()
				return nil
			}, re.SetModelAndConfigWithDefault(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }, nil))
			Expect(postJSON(e, "/", `{"model":"sentinel-model"}`).Code).To(Equal(200))
		}
	})

	It("exposes the diagnostic UUID before early errors and ignores client spoofing", func() {
		e := echo.New()
		e.POST("/", func(c echo.Context) error { return c.NoContent(204) },
			re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }))
		req := httptest.NewRequest("POST", "/", strings.NewReader(`{`))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("X-LocalAI-Diagnostic-ID", "spoof")
		rec := httptest.NewRecorder()
		e.ServeHTTP(rec, req)
		Expect(rec.Code).To(Equal(400))
		Expect(events).NotTo(BeEmpty())
		Expect(rec.Header().Get("X-LocalAI-Diagnostic-ID")).To(Equal(events[0].ID))
		Expect(events[0].ID).NotTo(Equal("spoof"))
	})

	It("preserves anonymous default, JSON, bearer and query selection with timing off and on", func() {
		for _, tc := range []struct {
			body, query, bearer string
			listing, lookup     int
		}{
			{`{}`, "", "", 1, 0}, {`{"model":"sentinel-model","prompt":"sentinel-prompt"}`, "", "", 1, 0},
			{`{}`, "", "sentinel-model", 0, 1}, {`{}`, "?model=sentinel-model", "sentinel-credential", 0, 0},
		} {
			var previous string
			recorder := ac.DiagnosticsRecorder
			for _, enabled := range []bool{false, true} {
				events = nil
				ac.DiagnosticsRecorder = nil
				if enabled {
					ac.DiagnosticsRecorder = recorder
				}
				e := echo.New()
				e.POST("/sentinel-path", func(c echo.Context) error {
					Expect(len(ends(diagnostics.PhaseExtraction))).To(Equal(map[bool]int{false: 0, true: 2}[enabled]))
					input := c.Get(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenAIRequest)
					Expect(re.SetOpenAIRequest(c)).To(Succeed())
					defer input.Cancel()
					Expect(diagnostics.Enabled(input.Context)).To(Equal(enabled))
					diagnostics.Mark(input.Context, diagnostics.PhaseModelInit)
					return c.String(200, input.Model)
				}, re.BuildFilteredFirstAvailableDefaultModel(nil), re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }))
				req := httptest.NewRequest(http.MethodPost, "/sentinel-path"+tc.query, strings.NewReader(tc.body))
				req.Header.Set("Content-Type", "application/json")
				req.Header.Set("X-Correlation-ID", "sentinel-correlation")
				if tc.bearer != "" {
					req.Header.Set("Authorization", "Bearer "+tc.bearer)
				}
				rec := httptest.NewRecorder()
				e.ServeHTTP(rec, req)
				Expect(rec.Code).To(Equal(200))
				Expect(rec.Body.String()).To(Equal("sentinel-model"))
				if !enabled {
					previous = rec.Body.String()
					Expect(events).To(BeEmpty())
					continue
				}
				Expect(rec.Body.String()).To(Equal(previous))
				Expect(ends(diagnostics.PhaseDefaultListing)).To(HaveLen(tc.listing))
				Expect(ends(diagnostics.PhaseBearerLookup)).To(HaveLen(tc.lookup))
				Expect(ends(diagnostics.PhaseBodyLookup)).To(HaveLen(1))
				Expect(ends(diagnostics.PhaseConfigLoadDefaults)).To(HaveLen(1))
				Expect(ends(diagnostics.PhaseConfigFilter)).NotTo(BeEmpty())
				id := events[0].ID
				_, err := uuid.Parse(id)
				Expect(err).NotTo(HaveOccurred())
				for _, event := range events {
					Expect(event.ID).To(Equal(id))
					Expect(event.Kind).To(Equal(diagnostics.KindRequest))
				}
				Expect(fmt.Sprint(events)).NotTo(ContainSubstring("sentinel"))
			}
			ac.DiagnosticsRecorder = recorder
		}
	})
	It("records bind, nil initializer result, missing, disabled and alias failures even when JSON returns nil", func() {
		disabled := true
		d := config.ModelConfig{Name: "disabled", Disabled: &disabled}
		d.Model = "disabled"
		a := config.ModelConfig{Name: "alias", Alias: "absent"}
		bcl.ReplaceModelConfigs([]config.ModelConfig{d, a})
		for _, tc := range []struct {
			body     string
			nilInput bool
			code     int
		}{
			{`{`, false, 400}, {`{}`, true, 400}, {`{"model":"missing"}`, false, 404}, {`{"model":"disabled"}`, false, 403}, {`{"model":"alias"}`, false, 400},
		} {
			events = nil
			e := echo.New()
			e.POST("/", func(c echo.Context) error { Fail("unexpected inference"); return nil }, re.SetModelAndConfig(func() schema.LocalAIRequest {
				if tc.nilInput {
					return nil
				}
				return new(schema.OpenAIRequest)
			}))
			rec := postJSON(e, "/", tc.body)
			Expect(rec.Code).To(Equal(tc.code))
			Expect(ends(diagnostics.PhaseExtraction)).To(HaveLen(1))
			Expect(ends(diagnostics.PhaseExtraction)[0].Outcome).To(Equal(diagnostics.OutcomeError))
			if strings.Contains(tc.body, "alias") {
				Expect(ends(diagnostics.PhaseAliasResolution)[0].Outcome).To(Equal(diagnostics.OutcomeError))
			}
		}
	})
	It("retains one identity across real failover attempts without default middleware", func() {
		primary := config.ModelConfig{Name: "sentinel-primary"}
		primary.Model = primary.Name
		fallback := config.ModelConfig{Name: "sentinel-fallback"}
		fallback.Model = fallback.Name
		chain := config.ModelConfig{Name: "sentinel-chain", Failover: &config.FailoverConfig{
			Targets: []config.FailoverTarget{{Model: primary.Name}, {Model: fallback.Name}},
		}}
		bcl.ReplaceModelConfigs([]config.ModelConfig{primary, fallback, chain})
		recorder := ac.DiagnosticsRecorder
		for _, enabled := range []bool{false, true} {
			events = nil
			ac.DiagnosticsRecorder = nil
			if enabled {
				ac.DiagnosticsRecorder = recorder
			}
			manager := failover.New(bcl)
			manager.Sync()
			re.SetFailoverManager(manager)
			var attempts []string
			e := echo.New()
			e.POST("/", func(c echo.Context) error {
				cfg := c.Get(CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig)
				attempts = append(attempts, cfg.Name)
				input := c.Get(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenAIRequest)
				Expect(input.Model).To(Equal(chain.Name))
				Expect(input.Prompt).To(Equal("sentinel-prompt"))
				if enabled {
					Expect(ends(diagnostics.PhaseExtraction)).To(HaveLen(len(attempts)))
					last := events[len(events)-1]
					Expect(last.Phase).To(Equal(diagnostics.PhaseExtraction))
					Expect(last.State).To(Equal(diagnostics.StateEnd))
					Expect(last.Outcome).To(Equal(diagnostics.OutcomeOK))
				}
				diagnostics.Mark(c.Request().Context(), diagnostics.PhaseModelInit)
				if cfg.Name == primary.Name {
					input.Prompt = "mutated"
					return echo.ErrBadGateway
				}
				return c.String(http.StatusOK, cfg.Name)
			}, re.SetModelAndConfigWithDefault(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }, nil))
			response := postJSON(e, "/", `{"model":"sentinel-chain","prompt":"sentinel-prompt"}`)
			Expect(attempts).To(Equal([]string{primary.Name, fallback.Name}))
			Expect(response.Code).To(Equal(http.StatusOK))
			Expect(response.Body.String()).To(Equal(fallback.Name))
			Expect(response.Header().Get(HeaderServedModel)).To(Equal(fallback.Name))
			Expect(response.Header().Get(HeaderFailover)).To(Equal("fallback"))
			if !enabled {
				Expect(events).To(BeEmpty())
				Expect(response.Header().Get(diagnostics.HeaderDiagnosticID)).To(BeEmpty())
				continue
			}
			Expect(ends(diagnostics.PhaseExtraction)).To(HaveLen(2))
			id := events[0].ID
			Expect(response.Header().Get(diagnostics.HeaderDiagnosticID)).To(Equal(id))
			_, err := uuid.Parse(id)
			Expect(err).NotTo(HaveOccurred())
			for _, event := range events {
				Expect(event.ID).To(Equal(id))
				Expect(event.Kind).To(Equal(diagnostics.KindRequest))
			}
			Expect(fmt.Sprint(events)).NotTo(ContainSubstring("sentinel"))
		}
	})
	It("retains identity across repeated segments and records cancellation before downstream", func() {
		ctx, cancel := context.WithCancel(context.Background())
		cancel()
		c := echo.New().NewContext(httptest.NewRequest("POST", "/", nil).WithContext(ctx), httptest.NewRecorder())
		h := re.BuildConstantDefaultModelNameMiddleware("sentinel-model")(func(c echo.Context) error {
			Expect(ends(diagnostics.PhaseExtraction)).NotTo(BeEmpty())
			return echo.ErrBadGateway
		})
		Expect(h(c)).To(Equal(echo.ErrBadGateway))
		Expect(h(c)).To(Equal(echo.ErrBadGateway))
		Expect(ends(diagnostics.PhaseExtraction)).To(HaveLen(2))
		for _, e := range events {
			Expect(e.ID).To(Equal(events[0].ID))
			if e.State == diagnostics.StateEnd {
				Expect(e.Outcome).To(Equal(diagnostics.OutcomeCanceled))
			}
		}
		Expect(c.Request().Context().Err()).To(Equal(context.Canceled))
	})
	It("inherits observation and request cancellation in Responses context rebuilding", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		c := echo.New().NewContext(httptest.NewRequest("POST", "/", nil).WithContext(ctx), httptest.NewRecorder())
		Expect(re.BuildConstantDefaultModelNameMiddleware("sentinel-model")(func(echo.Context) error { return nil })(c)).To(Succeed())
		input := &schema.OpenResponsesRequest{Model: "sentinel-model"}
		c.Set(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, input)
		cfg := config.ModelConfig{Name: "sentinel-model"}
		cfg.Model = "sentinel-model"
		c.Set(CONTEXT_LOCALS_KEY_MODEL_CONFIG, &cfg)
		Expect(re.SetOpenResponsesRequest(c)).To(Succeed())
		defer input.Cancel()
		Expect(diagnostics.Enabled(input.Context)).To(BeTrue())
		diagnostics.Mark(input.Context, diagnostics.PhaseModelInit)
		Expect(events[len(events)-1].ID).To(Equal(events[0].ID))
		cancel()
		Eventually(input.Context.Done()).Should(BeClosed())
	})
	It("isolates concurrent requests and a background operation on the same recorder", func() {
		var mu sync.Mutex
		ac.DiagnosticsRecorder = diagnostics.NewRecorder(func(e diagnostics.Event) { mu.Lock(); defer mu.Unlock(); events = append(events, e) })
		e := echo.New()
		e.POST("/", func(c echo.Context) error {
			diagnostics.Mark(c.Request().Context(), diagnostics.PhaseModelInit)
			return c.NoContent(204)
		}, re.BuildFilteredFirstAvailableDefaultModel(nil))
		const n = 12
		var wg sync.WaitGroup
		for i := 0; i < n; i++ {
			wg.Add(1)
			go func() { defer GinkgoRecover(); defer wg.Done(); Expect(postJSON(e, "/", `{}`).Code).To(Equal(204)) }()
		}
		reload := ac.DiagnosticsRecorder.Reload(context.Background())
		diagnostics.Mark(reload, diagnostics.PhaseReloadEnumeration)
		wg.Wait()
		groups := map[string][]diagnostics.Event{}
		for _, event := range events {
			groups[event.ID] = append(groups[event.ID], event)
		}
		Expect(groups).To(HaveLen(n + 1))
		requests := 0
		for _, group := range groups {
			if group[0].Kind == diagnostics.KindConfigReload {
				Expect(group).To(HaveLen(1))
				continue
			}
			requests++
			Expect(group[0].Phase).To(Equal(diagnostics.PhaseExtraction))
			Expect(group[0].State).To(Equal(diagnostics.StateStart))
			Expect(group[len(group)-2].Phase).To(Equal(diagnostics.PhaseExtraction))
			Expect(group[len(group)-2].State).To(Equal(diagnostics.StateEnd))
			Expect(group[len(group)-1].Phase).To(Equal(diagnostics.PhaseModelInit))
		}
		Expect(requests).To(Equal(n))
	})
	It("ends an empty default listing before downstream and preserves successful alias selection", func() {
		bcl.ReplaceModelConfigs(nil)
		c := echo.New().NewContext(httptest.NewRequest("POST", "/", nil), httptest.NewRecorder())
		Expect(re.BuildFilteredFirstAvailableDefaultModel(nil)(func(c echo.Context) error {
			Expect(ends(diagnostics.PhaseDefaultListing)).To(HaveLen(1))
			Expect(ends(diagnostics.PhaseExtraction)).To(HaveLen(1))
			return nil
		})(c)).To(Succeed())
		target := config.ModelConfig{Name: "target"}
		target.Model = "target"
		alias := config.ModelConfig{Name: "alias", Alias: "target"}
		bcl.ReplaceModelConfigs([]config.ModelConfig{target, alias})
		events = nil
		e := echo.New()
		e.POST("/", func(c echo.Context) error {
			Expect(c.Get(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenAIRequest).Model).To(Equal("alias"))
			Expect(c.Get(CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig).Name).To(Equal("target"))
			Expect(ends(diagnostics.PhaseAliasResolution)).To(HaveLen(1))
			Expect(ends(diagnostics.PhaseAliasResolution)[0].Outcome).To(Equal(diagnostics.OutcomeOK))
			return c.NoContent(204)
		}, re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }))
		Expect(postJSON(e, "/", `{"model":"alias"}`).Code).To(Equal(204))
	})

	It("records direct OpenAI and Responses extraction errors", func() {
		for _, setter := range []func(echo.Context) error{re.SetOpenAIRequest, re.SetOpenResponsesRequest} {
			events = nil
			c := echo.New().NewContext(httptest.NewRequest("POST", "/", nil), httptest.NewRecorder())
			Expect(setter(c)).To(Equal(echo.ErrBadRequest))
			Expect(ends(diagnostics.PhaseExtraction)).To(HaveLen(1))
			Expect(ends(diagnostics.PhaseExtraction)[0].Outcome).To(Equal(diagnostics.OutcomeError))
		}
	})

})