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>

backend/python/common/model_identity.py

@@ -6,7 +6,8 @@ model's backend, and the controller's health probe checks liveness rather than
 identity, so the request is dispatched to whatever now occupies the port and
 the caller gets a silent wrong-model answer (#10952).
 
-PredictOptions.ModelIdentity carries the model the request is for, so the
+Every request message that reaches a backend through the distributed router
+carries a ModelIdentity field naming the model the request is for, so the
 backend can reject it at the point of use. This module enforces that for every
 Python backend at once: all 36 of them build their server through
 grpc_auth.get_auth_interceptors(), so wiring it there needs no per-backend
@@ -31,16 +32,51 @@ MODEL_MISMATCH_SENTINEL = "model identity mismatch"
 
 _LOAD_METHOD = "/backend.Backend/LoadModel"
 
-# The four RPCs that carry PredictOptions. Nothing else has an identity field.
-# TTS and SoundGeneration are excluded on purpose: their `model` field is
-# already rewritten to a worker-local path by the controller's
-# FileStagingClient, so comparing it would reject valid requests.
+# Every RPC whose request message carries a ModelIdentity field. This set IS
+# the enforcement surface for all 36 Python backends: an RPC missing here is
+# silently unprotected, so model_identity_test.py pins the full list.
+#
+# The guard reads request.ModelIdentity generically, so nothing here is
+# modality-specific — a backend that does not implement an RPC simply never
+# sees it.
+#
+# TTS and SoundGeneration are guarded on ModelIdentity, NOT on their `model`
+# field: the controller's FileStagingClient rewrites `model` to a worker-local
+# absolute path, so comparing that would reject valid requests in distributed
+# mode. ModelIdentity is a separate, untranslated field for exactly that reason.
+#
+# AudioEncode/AudioDecode are absent deliberately: the opus codec backend they
+# target is loaded from a literal rather than a ModelConfig, so no value carries
+# the load-time/request-time equality guarantee this comparison depends on.
 _GUARDED_METHODS = frozenset(
     (
+        # PredictOptions RPCs (#10970)
         "/backend.Backend/Predict",
         "/backend.Backend/PredictStream",
         "/backend.Backend/Embedding",
         "/backend.Backend/TokenizeString",
+        # Remaining modalities
+        "/backend.Backend/GenerateImage",
+        "/backend.Backend/GenerateVideo",
+        "/backend.Backend/TTS",
+        "/backend.Backend/TTSStream",
+        "/backend.Backend/SoundGeneration",
+        "/backend.Backend/AudioTranscription",
+        "/backend.Backend/AudioTranscriptionStream",
+        "/backend.Backend/Detect",
+        "/backend.Backend/Depth",
+        "/backend.Backend/FaceVerify",
+        "/backend.Backend/FaceAnalyze",
+        "/backend.Backend/VoiceVerify",
+        "/backend.Backend/VoiceAnalyze",
+        "/backend.Backend/VoiceEmbed",
+        "/backend.Backend/Rerank",
+        "/backend.Backend/TokenClassify",
+        "/backend.Backend/Score",
+        "/backend.Backend/VAD",
+        "/backend.Backend/Diarize",
+        "/backend.Backend/SoundDetection",
+        "/backend.Backend/AudioTransform",
     )
 )
 
@@ -87,8 +123,10 @@ class ModelIdentityState:
 def _rebuild(handler, behavior):
     """Return a copy of `handler` with its behavior replaced.
 
-    Only unary-request handlers are ever passed here: LoadModel and the four
-    guarded RPCs all take a single request message.
+    Only unary-request handlers are ever passed here: LoadModel and every
+    entry in _GUARDED_METHODS take a single request message. The bidirectional
+    streams (AudioTranscriptionLive, AudioTransformStream, AudioToAudioStream,
+    Forward) are not guarded and never reach this function.
     """
     if handler.response_streaming:
         return grpc.unary_stream_rpc_method_handler(

backend/python/common/model_identity_test.py

@@ -214,9 +214,95 @@ class TestInterceptorBehavior(unittest.TestCase):
 
     def test_unguarded_rpcs_pass_through_untouched(self):
         handler = _handler(lambda request, context: "served")
-        wrapped = self._intercept("/backend.Backend/TTS", handler)
+        wrapped = self._intercept("/backend.Backend/Health", handler)
         self.assertIs(wrapped, handler)
 
 
+# Every modality request message now carries a ModelIdentity field, so every
+# modality RPC shares the guard. The set below is the enforcement surface for
+# all 36 Python backends at once: an RPC missing from it is silently
+# unprotected, which is the failure mode this class exists to catch.
+_EXPECTED_MODALITY_METHODS = (
+    "/backend.Backend/GenerateImage",
+    "/backend.Backend/GenerateVideo",
+    "/backend.Backend/TTS",
+    "/backend.Backend/TTSStream",
+    "/backend.Backend/SoundGeneration",
+    "/backend.Backend/AudioTranscription",
+    "/backend.Backend/AudioTranscriptionStream",
+    "/backend.Backend/Detect",
+    "/backend.Backend/Depth",
+    "/backend.Backend/FaceVerify",
+    "/backend.Backend/FaceAnalyze",
+    "/backend.Backend/VoiceVerify",
+    "/backend.Backend/VoiceAnalyze",
+    "/backend.Backend/VoiceEmbed",
+    "/backend.Backend/Rerank",
+    "/backend.Backend/TokenClassify",
+    "/backend.Backend/Score",
+    "/backend.Backend/VAD",
+    "/backend.Backend/Diarize",
+    "/backend.Backend/SoundDetection",
+    "/backend.Backend/AudioTransform",
+)
+
+
+class TestModalityMethods(unittest.TestCase):
+    def setUp(self):
+        self.interceptor = model_identity.ModelIdentityInterceptor()
+        self.interceptor.state.record("a.gguf")
+        self.served = []
+
+    def _call(self, method, identity, response_streaming=False):
+        def behavior(request, context):
+            self.served.append(method)
+            return "served"
+
+        handler = _handler(behavior, response_streaming=response_streaming)
+        wrapped = self.interceptor.intercept_service(
+            lambda _: handler, _FakeCallDetails(method)
+        )
+        fn = wrapped.unary_stream if response_streaming else wrapped.unary_unary
+        return fn(_Request(ModelIdentity=identity), _FakeContext())
+
+    def test_every_modality_rpc_is_guarded(self):
+        for method in _EXPECTED_MODALITY_METHODS:
+            with self.subTest(method=method):
+                self.assertIn(
+                    method,
+                    model_identity._GUARDED_METHODS,
+                    "{} is unprotected on all Python backends".format(method),
+                )
+
+    def test_every_modality_rpc_rejects_a_mismatch(self):
+        for method in _EXPECTED_MODALITY_METHODS:
+            streaming = method.endswith("Stream")
+            with self.subTest(method=method):
+                with self.assertRaises(_Aborted):
+                    self._call(method, "b.gguf", response_streaming=streaming)
+        self.assertEqual(self.served, [], "no request may reach the model")
+
+    def test_every_modality_rpc_serves_a_match(self):
+        for method in _EXPECTED_MODALITY_METHODS:
+            streaming = method.endswith("Stream")
+            self._call(method, "a.gguf", response_streaming=streaming)
+        self.assertEqual(len(self.served), len(_EXPECTED_MODALITY_METHODS))
+
+    # Compatibility: an old controller sends nothing, and the e2e backend suite
+    # drives real backends with bare request structs.
+    def test_every_modality_rpc_serves_without_an_identity(self):
+        for method in _EXPECTED_MODALITY_METHODS:
+            streaming = method.endswith("Stream")
+            self._call(method, "", response_streaming=streaming)
+        self.assertEqual(len(self.served), len(_EXPECTED_MODALITY_METHODS))
+
+    # AudioEncode/AudioDecode stay out: the opus codec backend is loaded from a
+    # literal, not a ModelConfig, so no value carries the structural guarantee
+    # the comparison depends on.
+    def test_codec_rpcs_stay_unguarded(self):
+        for method in ("/backend.Backend/AudioEncode", "/backend.Backend/AudioDecode"):
+            self.assertNotIn(method, model_identity._GUARDED_METHODS)
+
+
 if __name__ == "__main__":
     unittest.main()