[Bugfix] Reject /v1/completions for thinker+talker models when --trust-remote-code is unset (#5525)

Signed-off-by: ChoHee15 <cc5281@126.com>

tests/engine/test_single_stage_mode.py

@@ -12,6 +12,7 @@ import pytest
 from pytest_mock import MockerFixture
 from vllm.v1.engine.utils import EngineZmqAddresses
 
+from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
@@ -476,6 +477,89 @@ class TestSingleStageModeDetection:
         assert not hasattr(engine, "_omni_master_server")
 
 
+# ---------------------------------------------------------------------------
+# Endpoint-restriction resolution vs. tri-state trust_remote_code (#5495)
+# ---------------------------------------------------------------------------
+
+
+class TestEndpointRestrictionsTrustRemoteCode:
+    """Regression tests for issue #5495.
+
+    ``AsyncOmniEngine.__init__`` takes ``trust_remote_code: bool | None`` where
+    ``None`` means "not specified" (the CLI default, since ``--trust-remote-code``
+    is ``store_true`` and its absence is mapped to ``None``). Endpoint-restriction
+    resolution eagerly loads the HF config via vLLM's ``get_config``, which does
+    ``trust_remote_code |= ...`` internally and raises ``TypeError`` on ``None``.
+    That used to be swallowed, leaving ``endpoint_restrictions`` empty so
+    ``/v1/completions`` was never shut down and the request crashed stage-1.
+
+    ``__init__`` must collapse the ``None`` "not specified" case to ``False`` at
+    the ``get_pipeline_config`` call site so the restriction is still computed.
+    """
+
+    @pytest.fixture(autouse=True)
+    def clear_config_factory_caches(self):
+        """Isolate the process-wide ``functools.cache`` on StageConfigFactory.
+
+        ``get_hf_config`` / ``try_infer_model_type`` cache per-model results for
+        the whole process. These tests patch ``get_config`` with a synthetic
+        Qwen3-Omni config for ``"fake-model"``; without clearing on teardown the
+        synthetic entry would leak to later tests in the same worker (after the
+        mock is restored), making the suite order-dependent. Clear both before
+        and after.
+        """
+        StageConfigFactory.get_hf_config.cache_clear()
+        StageConfigFactory.try_infer_model_type.cache_clear()
+        yield
+        StageConfigFactory.get_hf_config.cache_clear()
+        StageConfigFactory.try_infer_model_type.cache_clear()
+
+    def _make_engine_no_thread(self, mocker: MockerFixture, **kwargs: Any) -> AsyncOmniEngine:
+        mocker.patch.object(
+            AsyncOmniEngine,
+            "_resolve_stage_configs",
+            return_value=("/fake/path", [_make_stage_cfg(0)]),
+        )
+        mocker.patch.object(AsyncOmniEngine, "_bootstrap_orchestrator")
+        mocker.patch("threading.Thread")
+        mocker.patch("concurrent.futures.Future")
+        return AsyncOmniEngine(model="fake-model", **kwargs)
+
+    @staticmethod
+    def _install_qwen3_omni_config(mocker: MockerFixture) -> None:
+        """Patch get_config to a thinker+talker config, mirroring vLLM's
+        ``trust_remote_code |= ...`` (raises TypeError on None)."""
+        from transformers import Qwen3OmniMoeConfig
+
+        def fake_get_config(model, trust_remote_code, **_):
+            trust_remote_code |= False  # TypeError if None reaches here
+            return Qwen3OmniMoeConfig(enable_audio_output=True)
+
+        mocker.patch("vllm_omni.config.config_factory.get_config", side_effect=fake_get_config)
+
+    @pytest.mark.parametrize("trc_kwargs", [{}, {"trust_remote_code": False}, {"trust_remote_code": True}])
+    def test_completions_restriction_survives(self, mocker: MockerFixture, trc_kwargs: dict):
+        """COMPLETIONS restriction is computed regardless of how (or whether)
+        trust_remote_code is passed — including the unset case that resolves to
+        the engine's ``None`` default."""
+        from vllm_omni.config.endpoint_policy import OmniServingCapability
+
+        self._install_qwen3_omni_config(mocker)
+        engine = self._make_engine_no_thread(mocker, **trc_kwargs)
+        assert any(r.capability is OmniServingCapability.COMPLETIONS for r in engine.endpoint_restrictions), (
+            f"COMPLETIONS restriction dropped for {trc_kwargs or 'unset (None default)'}"
+        )
+
+    def test_unset_trust_remote_code_defaults_to_none(self, mocker: MockerFixture):
+        """Guard the tri-state premise: the engine keeps ``None`` (not ``False``)
+        when trust_remote_code is not passed, so the call-site coercion is what
+        prevents the crash."""
+        import inspect
+
+        default = inspect.signature(AsyncOmniEngine.__init__).parameters["trust_remote_code"].default
+        assert default is None
+
+
 # ---------------------------------------------------------------------------
 # AsyncOmniEngine single-stage initialization paths
 # ---------------------------------------------------------------------------

tests/utils/test_tracking_parser.py

@@ -219,6 +219,46 @@ def test_store_true_explicit():
     assert ns.verbose is True
 
 
+def test_unset_trust_remote_code_falls_back_to_engine_default_none():
+    """Regression guard for issue #5495 (the CLI-plumbing half).
+
+    The CI test ``test_completions_rejected_for_thinker_talker`` does NOT pass
+    ``--trust-remote-code``. Because it is a ``store_true`` flag, an omitted
+    value is dropped from ``get_explicit_kwargs_dict()`` (it is not user-typed),
+    so it is never forwarded to ``AsyncOmni(**kwargs)`` and the engine falls back
+    to ``AsyncOmniEngine.__init__``'s default of ``None`` — NOT argparse's
+    ``False``. That ``None`` is what used to break endpoint-restriction
+    resolution (see ``TestQwen3OmniPipeline.
+    test_endpoint_restrictions_survive_unset_trust_remote_code``).
+
+    This test locks in the "CI does not pass the flag => effective value is
+    None" premise so a future refactor cannot silently change it.
+    """
+    import inspect
+
+    from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+
+    p = TrackingArgumentParser()
+    # Mirror how --trust-remote-code is registered upstream (store_true).
+    p.add_argument("--trust-remote-code", action="store_true")
+
+    # Scenario A: CI — flag omitted.
+    ns = p.parse_args([])
+    assert "trust_remote_code" not in ns.explicit_keys
+    kwargs = ns.get_explicit_kwargs_dict()
+    assert "trust_remote_code" not in kwargs
+    engine_default = inspect.signature(AsyncOmniEngine.__init__).parameters["trust_remote_code"].default
+    assert engine_default is None
+    effective = kwargs.get("trust_remote_code", engine_default)
+    assert effective is None  # <- the value that reached the buggy restriction path
+
+    # Scenario B: flag passed (what earlier repros used) => True, restrictions fine.
+    ns2 = p.parse_args(["--trust-remote-code"])
+    assert "trust_remote_code" in ns2.explicit_keys
+    kwargs2 = ns2.get_explicit_kwargs_dict()
+    assert kwargs2.get("trust_remote_code", engine_default) is True
+
+
 def test_store_false_default():
     """Ensure that store false is handled correctly when omitted."""
     p = TrackingArgumentParser()

vllm_omni/engine/async_omni_engine.py

@@ -213,6 +213,12 @@ class AsyncOmniEngine:
         # beforehand. The stage CLI exposes the same deploy YAML through
         # stage_configs_path.
         deploy_config_path = kwargs.get("deploy_config") or kwargs.get("stage_configs_path")
+        # ``trust_remote_code`` is tri-state (bool | None): ``None`` means "not
+        # specified" so stage-config resolution can defer to the deploy yaml's
+        # per-stage value (see ``with_trust_remote_code_override``). The
+        # restriction path below loads the top-level HF config via vLLM's
+        # ``get_config``, which needs a real bool, so collapse ``None`` to the
+        # default ``False`` here (#5495).
         pipeline_config = StageConfigFactory.get_pipeline_config(
             model=model,
             trust_remote_code=bool(trust_remote_code),