[Frontend] Share regular and duplex chat initialization (#7646)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>

tests/entrypoints/openai_api/test_api_server_guards.py

@@ -1053,7 +1053,8 @@ async def test_pure_diffusion_speech_forwards_media_access_args(monkeypatch, tmp
 
 
 @pytest.mark.asyncio
-async def test_multistage_app_state_key_snapshot(monkeypatch, tmp_path) -> None:
+@pytest.mark.parametrize("supported_tasks", [("generate",), ("embed",)])
+async def test_multistage_app_state_key_snapshot(monkeypatch, mocker, tmp_path, supported_tasks) -> None:
     """Lock multi-stage ``app.state`` keys after init, including live vs None.
 
     Fails if chat/speech/video/realtime/tokenization keys disappear or are
@@ -1068,6 +1069,7 @@ async def test_multistage_app_state_key_snapshot(monkeypatch, tmp_path) -> None:
             parallel_config=SimpleNamespace(_api_process_rank=0),
         ),
     )
+    mocker.patch.object(engine, "get_supported_tasks", return_value=supported_tasks)
 
     class _FakeModels:
         def __init__(self, *args, **kwargs):
@@ -1078,10 +1080,12 @@ async def test_multistage_app_state_key_snapshot(monkeypatch, tmp_path) -> None:
 
     class _FakeCtor:
         def __init__(self, *args, **kwargs):
-            pass
+            self.args = args
+            self.kwargs = kwargs
+            self.warmup_calls = 0
 
         def warmup(self):
-            return None
+            self.warmup_calls += 1
 
     deploy = tmp_path / "deploy.yaml"
     deploy.write_text("speech_cache:\n  resolve_max_bytes: 1024\n  speaker_max_bytes: 8\n")
@@ -1119,17 +1123,33 @@ async def test_multistage_app_state_key_snapshot(monkeypatch, tmp_path) -> None:
     monkeypatch.setattr(api_server, "OmniOpenAIServingVideo", _FakeCtor)
 
     state = State()
-    await api_server.omni_init_app_state(engine, state, _minimal_args())
+    await api_server.omni_init_app_state(engine, state, _minimal_args(log_error_stack=True))
     assert speech_kwargs["speech_cache_config"].resolve_max_bytes == 1024
     assert speech_kwargs["speech_cache_config"].resolve_max_entries == 2048
     assert speech_kwargs["speech_cache_config"].speaker_max_bytes == 8
 
+    disabled = (
+        set()
+        if "generate" in supported_tasks
+        else {"openai_serving_chat", "openai_serving_chat_batch", "openai_streaming_video"}
+    )
     _assert_app_state_snapshot(
         state,
         expected_keys=_MULTISTAGE_APP_STATE_KEYS,
-        must_be_wired=_MULTISTAGE_MUST_BE_WIRED,
-        must_be_none=_MULTISTAGE_MUST_BE_NONE,
+        must_be_wired=_MULTISTAGE_MUST_BE_WIRED - disabled,
+        must_be_none=_MULTISTAGE_MUST_BE_NONE | disabled,
     )
+    assert state.online_renderer.kwargs["log_error_stack"] is True
+    assert state.online_renderer.warmup_calls == 1
+    for service in (state.openai_serving_chat, state.openai_serving_chat_batch):
+        if "generate" in supported_tasks:
+            assert service.kwargs["online_renderer"] is state.online_renderer
+        else:
+            assert service is None
+    if "generate" in supported_tasks:
+        assert state.openai_serving_responses.args[2] is state.online_renderer
+    else:
+        assert state.openai_serving_responses is None
 
 
 @pytest.mark.parametrize("count", [None, 0, -1, "2", True])

tests/entrypoints/openai_api/test_duplex_api_server.py

@@ -166,6 +166,7 @@ def _minimal_args(**overrides) -> SimpleNamespace:
         enable_force_include_usage=False,
         enable_log_outputs=False,
         enable_log_deltas=False,
+        log_error_stack=False,
     )
     for key, value in overrides.items():
         setattr(args, key, value)
@@ -316,7 +317,8 @@ def test_minicpmo_serving_profiles(monkeypatch, deploy_name, expected):
 
 
 @pytest.mark.asyncio
-async def test_duplex_app_state_wires_only_the_session_surfaces(monkeypatch) -> None:
+@pytest.mark.parametrize("log_error_stack", [False, True])
+async def test_duplex_app_state_wires_only_the_session_surfaces(monkeypatch, mocker, log_error_stack) -> None:
     """Lock the duplex ``app.state``: the two session-backed surfaces are live, every turn-based service is None.
 
     Fails if a turn-based service is wired into a duplex server (its route
@@ -326,12 +328,12 @@ async def test_duplex_app_state_wires_only_the_session_surfaces(monkeypatch) ->
     the turn-based chat service.
     """
     monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels)
-    monkeypatch.setattr(api_server, "OnlineRenderer", lambda **kwargs: SimpleNamespace(**kwargs))
+    renderer_ctor = mocker.patch.object(api_server, "OnlineRenderer")
     monkeypatch.setattr(api_server, "OmniOpenAIServingChat", lambda **kwargs: SimpleNamespace(**kwargs))
     engine = _FakeDuplexOmni()
     state = State()
 
-    await api_server.omni_init_app_state(engine, state, _minimal_args())
+    await api_server.omni_init_app_state(engine, state, _minimal_args(log_error_stack=log_error_stack))
 
     present = {key for key in _DUPLEX_APP_STATE_KEYS if hasattr(state, key)}
     assert present == _DUPLEX_APP_STATE_KEYS
@@ -341,6 +343,10 @@ async def test_duplex_app_state_wires_only_the_session_surfaces(monkeypatch) ->
     assert not unexpectedly_set, f"turn-based services wired into a duplex server: {unexpectedly_set}"
     assert isinstance(state.openai_serving_duplex, OmniDuplexSessionHandler)
     assert state.openai_serving_chat is not None
+    renderer_ctor.assert_called_once()
+    assert renderer_ctor.call_args.kwargs["log_error_stack"] is log_error_stack
+    assert state.openai_serving_chat.online_renderer is state.online_renderer is renderer_ctor.return_value
+    state.online_renderer.warmup.assert_called_once_with()
     assert state.engine_client is engine
     assert state.vllm_config is engine._vllm_config
 
@@ -498,16 +504,27 @@ async def test_warmup_gate_lets_the_warmup_connection_and_plain_servers_through(
 
 
 @pytest.mark.asyncio
-async def test_a_model_that_does_not_declare_chat_completions_does_not_get_the_route(monkeypatch) -> None:
-    """The decision stays the model's: no capability, no chat service, and the route says so."""
+@pytest.mark.parametrize(
+    "supports_chat, supported_tasks", [(False, ("generate",)), (True, ())], ids=["no-capability", "no-generation"]
+)
+async def test_duplex_chat_initialization_respects_model_gates(
+    monkeypatch, mocker, supports_chat, supported_tasks
+) -> None:
+    """Neither gate may construct or warm up a chat renderer."""
     monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels)
+    renderer_ctor = mocker.patch.object(api_server, "OnlineRenderer")
+    chat_ctor = mocker.patch.object(api_server, "OmniOpenAIServingChat")
     engine = _FakeDuplexOmni()
-    engine._capabilities = DuplexCapabilities(supports_chat_completions=False)
+    engine._capabilities = DuplexCapabilities(supports_chat_completions=supports_chat)
+    mocker.patch.object(engine, "get_supported_tasks", return_value=supported_tasks)
     state = State()
 
     await api_server.omni_init_app_state(engine, state, _minimal_args())
 
     assert state.openai_serving_chat is None
+    assert state.openai_serving_chat_batch is None
+    renderer_ctor.assert_not_called()
+    chat_ctor.assert_not_called()
     assert isinstance(state.openai_serving_duplex, OmniDuplexSessionHandler)
 
 
@@ -535,3 +552,31 @@ def test_turn_deployment_rejects_explicit_duplex_without_falling_back_to_stt(fla
             with pytest.raises(WebSocketDisconnect) as exc:
                 websocket.receive_text()
             assert exc.value.code == 1008
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "explicit_template, tokenizer_template, tokenizer_error, expected",
+    [
+        ("explicit", None, False, "explicit"),
+        (None, "tokenizer", False, None),
+        (None, None, False, "model-json"),
+        (None, None, True, "model-json"),
+    ],
+)
+async def test_shared_chat_template_resolution(
+    mocker, explicit_template, tokenizer_template, tokenizer_error, expected
+):
+    engine = _FakeDuplexOmni()
+    tokenizer = mocker.patch.object(
+        engine,
+        "get_tokenizer",
+        return_value=SimpleNamespace(chat_template=tokenizer_template),
+        side_effect=RuntimeError("tokenizer unavailable") if tokenizer_error else None,
+    )
+    mocker.patch.object(api_server, "load_chat_template", return_value=explicit_template)
+    model_template = mocker.patch.object(api_server, "_load_model_chat_template_json", return_value="model-json")
+
+    assert await api_server._resolve_chat_template(engine, _minimal_args()) == expected
+    assert tokenizer.call_count == (0 if explicit_template is not None else 1)
+    assert model_template.call_count == (1 if expected == "model-json" else 0)