fix(distributed): reject wrong-model requests on the remaining modalities (#10990)

#10970 gave the four PredictOptions RPCs a model-identity check so a backend reached through a stale distributed route rejects the request instead of answering from whatever model it holds (#10952). Every other modality shares that exposure: the route is cached by host:port, a worker can recycle a stopped backend's port for another model's backend, and a liveness-only probe cannot tell a stale row from a valid one. Extends the same mechanism to the 21 remaining request messages that reach a backend through the router, using the pattern #10970 established rather than a parallel one: - proto: ModelIdentity on each modality request message. - controller: populated from ModelConfig.Model at the call site that also builds ModelOptions, so load-time and request-time values are equal by construction. - backends: one generic guard in pkg/grpc/server.go (27 Go backends), the method set in backend/python/common (36 Python backends), llama-cpp (AudioTranscription/Stream, Rerank, Score) and privacy-filter (TokenClassify). - reconcile already drops the stale row on IsModelMismatch; no change. TTSRequest and SoundGenerationRequest get a SEPARATE ModelIdentity field rather than reusing their existing `model`: FileStagingClient rewrites `model` to a worker-local path, so comparing it would reject valid requests in exactly the configuration this guards. AudioEncode/AudioDecode are deliberately left unguarded: the opus codec backend is loaded from a literal rather than a ModelConfig, so no value carries the equality guarantee the comparison depends on. The four bidirectional stream RPCs are out of scope; they bypass reconcile. Empty means skip on both sides, so an old controller, an old backend, and the bare request structs in tests/e2e-backends all keep working. Assisted-by: Claude Code:claude-opus-4-8 [Read] [Edit] [Bash] Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
2026-07-30 09:57:57 -04:00 · 2026-07-20 21:58:19 +02:00
parent a784cf669f
commit 1cd7d63c7b
29 changed files with 1104 additions and 85 deletions
--- a/pkg/grpc/model_identity_modalities_test.go
+++ b/pkg/grpc/model_identity_modalities_test.go
@@ -0,0 +1,226 @@
+package grpc
+
+import (
+	"context"
+
+	"github.com/mudler/LocalAI/pkg/grpc/base"
+	"github.com/mudler/LocalAI/pkg/grpc/grpcerrors"
+	pb "github.com/mudler/LocalAI/pkg/grpc/proto"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+// modalityBackend answers every non-PredictOptions inference RPC successfully
+// and counts the calls that reached it. `served` is therefore the signal for
+// "the identity guard let this through" — the same role it plays in
+// model_identity_test.go for the four PredictOptions RPCs.
+type modalityBackend struct {
+	base.SingleThread
+
+	loaded string
+	served int
+}
+
+func (b *modalityBackend) Load(opts *pb.ModelOptions) error {
+	b.loaded = opts.Model
+	return nil
+}
+
+func (b *modalityBackend) GenerateImage(*pb.GenerateImageRequest) error { b.served++; return nil }
+func (b *modalityBackend) GenerateVideo(*pb.GenerateVideoRequest) error { b.served++; return nil }
+func (b *modalityBackend) TTS(*pb.TTSRequest) error                     { b.served++; return nil }
+
+func (b *modalityBackend) TTSStream(_ *pb.TTSRequest, ch chan []byte) error {
+	b.served++
+	ch <- []byte{0}
+	close(ch)
+	return nil
+}
+
+func (b *modalityBackend) SoundGeneration(*pb.SoundGenerationRequest) error { b.served++; return nil }
+
+func (b *modalityBackend) AudioTranscription(context.Context, *pb.TranscriptRequest) (pb.TranscriptResult, error) {
+	b.served++
+	return pb.TranscriptResult{Text: "ok"}, nil
+}
+
+func (b *modalityBackend) AudioTranscriptionStream(_ context.Context, _ *pb.TranscriptRequest, ch chan *pb.TranscriptStreamResponse) error {
+	b.served++
+	ch <- &pb.TranscriptStreamResponse{Delta: "ok"}
+	close(ch)
+	return nil
+}
+
+func (b *modalityBackend) Detect(*pb.DetectOptions) (pb.DetectResponse, error) {
+	b.served++
+	return pb.DetectResponse{}, nil
+}
+
+func (b *modalityBackend) Depth(*pb.DepthRequest) (pb.DepthResponse, error) {
+	b.served++
+	return pb.DepthResponse{}, nil
+}
+
+func (b *modalityBackend) FaceVerify(*pb.FaceVerifyRequest) (pb.FaceVerifyResponse, error) {
+	b.served++
+	return pb.FaceVerifyResponse{}, nil
+}
+
+func (b *modalityBackend) FaceAnalyze(*pb.FaceAnalyzeRequest) (pb.FaceAnalyzeResponse, error) {
+	b.served++
+	return pb.FaceAnalyzeResponse{}, nil
+}
+
+func (b *modalityBackend) VoiceVerify(*pb.VoiceVerifyRequest) (pb.VoiceVerifyResponse, error) {
+	b.served++
+	return pb.VoiceVerifyResponse{}, nil
+}
+
+func (b *modalityBackend) VoiceAnalyze(*pb.VoiceAnalyzeRequest) (pb.VoiceAnalyzeResponse, error) {
+	b.served++
+	return pb.VoiceAnalyzeResponse{}, nil
+}
+
+func (b *modalityBackend) VoiceEmbed(*pb.VoiceEmbedRequest) (pb.VoiceEmbedResponse, error) {
+	b.served++
+	return pb.VoiceEmbedResponse{}, nil
+}
+
+func (b *modalityBackend) VAD(*pb.VADRequest) (pb.VADResponse, error) {
+	b.served++
+	return pb.VADResponse{}, nil
+}
+
+func (b *modalityBackend) Diarize(*pb.DiarizeRequest) (pb.DiarizeResponse, error) {
+	b.served++
+	return pb.DiarizeResponse{}, nil
+}
+
+func (b *modalityBackend) SoundDetection(context.Context, *pb.SoundDetectionRequest) (*pb.SoundDetectionResponse, error) {
+	b.served++
+	return &pb.SoundDetectionResponse{}, nil
+}
+
+func (b *modalityBackend) AudioTransform(*pb.AudioTransformRequest) (*pb.AudioTransformResult, error) {
+	b.served++
+	return &pb.AudioTransformResult{}, nil
+}
+
+var _ AIModel = (*modalityBackend)(nil)
+
+// callAllModalities exercises every remaining inference RPC that the generic Go
+// server implements, sending `identity` in each request's ModelIdentity field.
+// One call per RPC, so the returned error count is also the RPC count.
+//
+// Rerank, Score and TokenClassify are absent on purpose: the generic Go server
+// does not implement them (they fall through to UnimplementedBackendServer), so
+// only the C++ and Python backends can enforce them.
+func callAllModalities(c Backend, identity string) map[string]error {
+	ctx := context.Background()
+	errs := map[string]error{}
+
+	_, errs["GenerateImage"] = c.GenerateImage(ctx, &pb.GenerateImageRequest{ModelIdentity: identity})
+	_, errs["GenerateVideo"] = c.GenerateVideo(ctx, &pb.GenerateVideoRequest{ModelIdentity: identity})
+	_, errs["TTS"] = c.TTS(ctx, &pb.TTSRequest{ModelIdentity: identity})
+	errs["TTSStream"] = c.TTSStream(ctx, &pb.TTSRequest{ModelIdentity: identity}, func(*pb.Reply) {})
+	_, errs["SoundGeneration"] = c.SoundGeneration(ctx, &pb.SoundGenerationRequest{ModelIdentity: identity})
+	_, errs["AudioTranscription"] = c.AudioTranscription(ctx, &pb.TranscriptRequest{ModelIdentity: identity})
+	errs["AudioTranscriptionStream"] = c.AudioTranscriptionStream(ctx, &pb.TranscriptRequest{ModelIdentity: identity}, func(*pb.TranscriptStreamResponse) {})
+	_, errs["Detect"] = c.Detect(ctx, &pb.DetectOptions{ModelIdentity: identity})
+	_, errs["Depth"] = c.Depth(ctx, &pb.DepthRequest{ModelIdentity: identity})
+	_, errs["FaceVerify"] = c.FaceVerify(ctx, &pb.FaceVerifyRequest{ModelIdentity: identity})
+	_, errs["FaceAnalyze"] = c.FaceAnalyze(ctx, &pb.FaceAnalyzeRequest{ModelIdentity: identity})
+	_, errs["VoiceVerify"] = c.VoiceVerify(ctx, &pb.VoiceVerifyRequest{ModelIdentity: identity})
+	_, errs["VoiceAnalyze"] = c.VoiceAnalyze(ctx, &pb.VoiceAnalyzeRequest{ModelIdentity: identity})
+	_, errs["VoiceEmbed"] = c.VoiceEmbed(ctx, &pb.VoiceEmbedRequest{ModelIdentity: identity})
+	_, errs["VAD"] = c.VAD(ctx, &pb.VADRequest{ModelIdentity: identity})
+	_, errs["Diarize"] = c.Diarize(ctx, &pb.DiarizeRequest{ModelIdentity: identity})
+	_, errs["SoundDetection"] = c.SoundDetection(ctx, &pb.SoundDetectionRequest{ModelIdentity: identity})
+	_, errs["AudioTransform"] = c.AudioTransform(ctx, &pb.AudioTransformRequest{ModelIdentity: identity})
+
+	return errs
+}
+
+// #10970 protected only the four PredictOptions RPCs. Every other modality
+// shares the exposure — the controller's cached row names a host:port, the port
+// can be recycled by another model's backend, and a liveness-only probe cannot
+// tell the difference — so each one must reject a request naming another model
+// rather than answer from whatever it holds.
+var _ = Describe("modality RPC model identity guard", func() {
+	newServed := func(addr, loadedModel string) (Backend, *modalityBackend) {
+		b := &modalityBackend{}
+		Provide(addr, b)
+		c := NewClient(addr, true, nil, false)
+		_, err := c.LoadModel(context.Background(), &pb.ModelOptions{Model: loadedModel})
+		Expect(err).ToNot(HaveOccurred())
+		Expect(b.loaded).To(Equal(loadedModel))
+		return c, b
+	}
+
+	It("rejects every modality RPC when the identity names another model", func() {
+		c, b := newServed("test://modality-mismatch", "a.gguf")
+
+		for rpc, err := range callAllModalities(c, "b.gguf") {
+			Expect(err).To(HaveOccurred(), "%s must reject a wrong-model request", rpc)
+			Expect(grpcerrors.IsModelMismatch(err)).To(BeTrue(), "%s: want a mismatch error, got %v", rpc, err)
+			// The router reacts differently to the two signals, so a mismatch
+			// must never be mistaken for a not-loaded.
+			Expect(grpcerrors.IsModelNotLoaded(err)).To(BeFalse(), "%s", rpc)
+		}
+		Expect(b.served).To(Equal(0), "no request may reach the model on a mismatch")
+	})
+
+	It("serves when the identity matches the loaded model", func() {
+		c, b := newServed("test://modality-match", "a.gguf")
+
+		errs := callAllModalities(c, "a.gguf")
+		for rpc, err := range errs {
+			Expect(err).ToNot(HaveOccurred(), "%s", rpc)
+		}
+		Expect(b.served).To(Equal(len(errs)))
+	})
+
+	// Compatibility, old controller -> new backend. Every existing deployment
+	// sends no identity, and tests/e2e-backends/backend_test.go drives real
+	// backends with bare request structs at many call sites. Tightening this
+	// breaks all of them, so it must fail here first.
+	It("serves when the request carries no identity", func() {
+		c, b := newServed("test://modality-empty-request", "a.gguf")
+
+		errs := callAllModalities(c, "")
+		for rpc, err := range errs {
+			Expect(err).ToNot(HaveOccurred(), "%s", rpc)
+		}
+		Expect(b.served).To(Equal(len(errs)))
+	})
+
+	// The backend side of the same rule: a model loaded without an identity
+	// (an old controller did the load) cannot judge anything, so it must serve.
+	It("serves when the backend has no recorded identity", func() {
+		c, b := newServed("test://modality-empty-loaded", "")
+
+		errs := callAllModalities(c, "b.gguf")
+		for rpc, err := range errs {
+			Expect(err).ToNot(HaveOccurred(), "%s", rpc)
+		}
+		Expect(b.served).To(Equal(len(errs)))
+	})
+
+	// AudioEncode/AudioDecode are deliberately NOT guarded: the opus codec
+	// backend they target is loaded with a literal model.WithModel("opus") and
+	// has no ModelConfig, so there is no value with the structural guarantee
+	// the rest of this mechanism depends on. Sending an identity that merely
+	// looks right would be convention, not construction — the exact
+	// false-rejection risk #10970 refused to take. Pinned here so a future
+	// "completeness" pass has to make that decision deliberately.
+	It("leaves the codec RPCs unguarded", func() {
+		c, _ := newServed("test://modality-codec", "a.gguf")
+
+		// Whatever these answer, it must never be a model mismatch: no
+		// identity is carried, so nothing is compared.
+		_, err := c.AudioEncode(context.Background(), &pb.AudioEncodeRequest{})
+		Expect(grpcerrors.IsModelMismatch(err)).To(BeFalse())
+		_, err = c.AudioDecode(context.Background(), &pb.AudioDecodeRequest{})
+		Expect(grpcerrors.IsModelMismatch(err)).To(BeFalse())
+	})
+})
--- a/pkg/grpc/server.go
+++ b/pkg/grpc/server.go
@@ -34,13 +34,20 @@ type server struct {
 	pb.UnimplementedBackendServer
 	llm AIModel

-	// identityMu guards loadedIdentity: LoadModel writes it, the
-	// PredictOptions RPCs read it, and nothing serialises those against each
-	// other (llm.Locking() is the model's own lock, and it is optional).
+	// identityMu guards loadedIdentity: LoadModel writes it, the inference
+	// RPCs read it, and nothing serialises those against each other
+	// (llm.Locking() is the model's own lock, and it is optional).
 	identityMu     sync.RWMutex
 	loadedIdentity string
 }

+// identifiedRequest is satisfied by every request message that carries a
+// ModelIdentity field. protoc-gen-go generates GetModelIdentity with a
+// nil-receiver guard, so a typed-nil request is safe to pass here.
+type identifiedRequest interface {
+	GetModelIdentity() string
+}
+
 // checkModelIdentity reports an error when the request names a model other
 // than the one this process loaded. It is the point-of-use half of the fix for
 // #10952: in distributed mode a worker can recycle a stopped backend's gRPC
@@ -53,18 +60,23 @@ type server struct {
 // requests; the loaded side is empty when such a controller performed the
 // load. Neither can judge the other, and a false rejection is far worse than
 // the miss.
-func (s *server) checkModelIdentity(in *pb.PredictOptions) error {
-	if in == nil || in.ModelIdentity == "" {
+//
+// One guard covers every modality because they all share one exposure: the
+// route is cached by address, the address can be recycled, and only the
+// backend can tell. Which RPCs call it is the whole enforcement surface — see
+// pkg/grpc/model_identity_modalities_test.go, which drives all of them.
+func (s *server) checkModelIdentity(in identifiedRequest) error {
+	if in == nil || in.GetModelIdentity() == "" {
 		return nil
 	}
 	s.identityMu.RLock()
 	loaded := s.loadedIdentity
 	s.identityMu.RUnlock()

-	if loaded == "" || loaded == in.ModelIdentity {
+	if loaded == "" || loaded == in.GetModelIdentity() {
 		return nil
 	}
-	return grpcerrors.ModelMismatch("grpc-server", loaded, in.ModelIdentity)
+	return grpcerrors.ModelMismatch("grpc-server", loaded, in.GetModelIdentity())
 }

 func (s *server) Health(ctx context.Context, in *pb.HealthMessage) (*pb.Reply, error) {
@@ -124,6 +136,9 @@ func (s *server) Predict(ctx context.Context, in *pb.PredictOptions) (*pb.Reply,
 }

 func (s *server) GenerateImage(ctx context.Context, in *pb.GenerateImageRequest) (*pb.Result, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -136,6 +151,9 @@ func (s *server) GenerateImage(ctx context.Context, in *pb.GenerateImageRequest)
 }

 func (s *server) GenerateVideo(ctx context.Context, in *pb.GenerateVideoRequest) (*pb.Result, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -148,6 +166,9 @@ func (s *server) GenerateVideo(ctx context.Context, in *pb.GenerateVideoRequest)
 }

 func (s *server) TTS(ctx context.Context, in *pb.TTSRequest) (*pb.Result, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -160,6 +181,9 @@ func (s *server) TTS(ctx context.Context, in *pb.TTSRequest) (*pb.Result, error)
 }

 func (s *server) TTSStream(in *pb.TTSRequest, stream pb.Backend_TTSStreamServer) error {
+	if err := s.checkModelIdentity(in); err != nil {
+		return err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -181,6 +205,9 @@ func (s *server) TTSStream(in *pb.TTSRequest, stream pb.Backend_TTSStreamServer)
 }

 func (s *server) SoundGeneration(ctx context.Context, in *pb.SoundGenerationRequest) (*pb.Result, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -193,6 +220,9 @@ func (s *server) SoundGeneration(ctx context.Context, in *pb.SoundGenerationRequ
 }

 func (s *server) Detect(ctx context.Context, in *pb.DetectOptions) (*pb.DetectResponse, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -205,6 +235,9 @@ func (s *server) Detect(ctx context.Context, in *pb.DetectOptions) (*pb.DetectRe
 }

 func (s *server) Depth(ctx context.Context, in *pb.DepthRequest) (*pb.DepthResponse, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -217,6 +250,9 @@ func (s *server) Depth(ctx context.Context, in *pb.DepthRequest) (*pb.DepthRespo
 }

 func (s *server) FaceVerify(ctx context.Context, in *pb.FaceVerifyRequest) (*pb.FaceVerifyResponse, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -229,6 +265,9 @@ func (s *server) FaceVerify(ctx context.Context, in *pb.FaceVerifyRequest) (*pb.
 }

 func (s *server) FaceAnalyze(ctx context.Context, in *pb.FaceAnalyzeRequest) (*pb.FaceAnalyzeResponse, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -241,6 +280,9 @@ func (s *server) FaceAnalyze(ctx context.Context, in *pb.FaceAnalyzeRequest) (*p
 }

 func (s *server) VoiceVerify(ctx context.Context, in *pb.VoiceVerifyRequest) (*pb.VoiceVerifyResponse, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -253,6 +295,9 @@ func (s *server) VoiceVerify(ctx context.Context, in *pb.VoiceVerifyRequest) (*p
 }

 func (s *server) VoiceAnalyze(ctx context.Context, in *pb.VoiceAnalyzeRequest) (*pb.VoiceAnalyzeResponse, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -265,6 +310,9 @@ func (s *server) VoiceAnalyze(ctx context.Context, in *pb.VoiceAnalyzeRequest) (
 }

 func (s *server) VoiceEmbed(ctx context.Context, in *pb.VoiceEmbedRequest) (*pb.VoiceEmbedResponse, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -277,6 +325,9 @@ func (s *server) VoiceEmbed(ctx context.Context, in *pb.VoiceEmbedRequest) (*pb.
 }

 func (s *server) AudioTranscription(ctx context.Context, in *pb.TranscriptRequest) (*pb.TranscriptResult, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -319,6 +370,9 @@ func (s *server) AudioTranscription(ctx context.Context, in *pb.TranscriptReques
 }

 func (s *server) AudioTranscriptionStream(in *pb.TranscriptRequest, stream pb.Backend_AudioTranscriptionStreamServer) error {
+	if err := s.checkModelIdentity(in); err != nil {
+		return err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -536,6 +590,9 @@ func (s *server) StoresFind(ctx context.Context, in *pb.StoresFindOptions) (*pb.
 }

 func (s *server) VAD(ctx context.Context, in *pb.VADRequest) (*pb.VADResponse, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -548,6 +605,9 @@ func (s *server) VAD(ctx context.Context, in *pb.VADRequest) (*pb.VADResponse, e
 }

 func (s *server) Diarize(ctx context.Context, in *pb.DiarizeRequest) (*pb.DiarizeResponse, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -560,6 +620,9 @@ func (s *server) Diarize(ctx context.Context, in *pb.DiarizeRequest) (*pb.Diariz
 }

 func (s *server) SoundDetection(ctx context.Context, in *pb.SoundDetectionRequest) (*pb.SoundDetectionResponse, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()
@@ -592,6 +655,9 @@ func (s *server) AudioDecode(ctx context.Context, in *pb.AudioDecodeRequest) (*p
 }

 func (s *server) AudioTransform(ctx context.Context, in *pb.AudioTransformRequest) (*pb.AudioTransformResult, error) {
+	if err := s.checkModelIdentity(in); err != nil {
+		return nil, err
+	}
 	if s.llm.Locking() {
 		s.llm.Lock()
 		defer s.llm.Unlock()