perf(vllm): derive required prefix token IDs in Gym

responses_api_models/vllm_model/app.py

@@ -697,6 +697,8 @@ def _preprocess_chat_completion_create_params(self, request: Request, body_dict:
                 # No user message found — create one with just the audio blocks.
                 body_dict.setdefault("messages", []).append({"role": "user", "content": list(audio_blocks)})
 
+        if self.config.return_token_id_information:
+            self._derive_required_prefix_token_ids(body_dict)
         self._apply_sampling_overrides(body_dict)
         self._validate_single_choice_token_request(body_dict)
         if self._external_capture_enabled:
@@ -1251,6 +1253,23 @@ def _extract_vllm_response_token_ids(
             cls._require_token_id_list(generation_value, "choice.token_ids"),
         )
 
+    @classmethod
+    def _derive_required_prefix_token_ids(cls, body_dict: Dict[str, Any]) -> None:
+        if body_dict.get("required_prefix_token_ids") is not None:
+            return
+
+        for message in reversed(body_dict.get("messages", [])):
+            if not isinstance(message, dict) or message.get("role") != "assistant":
+                continue
+            token_bundle = cls._extract_message_token_bundle(message)
+            if token_bundle is None:
+                continue
+            body_dict["required_prefix_token_ids"] = [
+                *token_bundle["prompt_token_ids"],
+                *token_bundle["generation_token_ids"],
+            ]
+            return
+
     def _extract_choice_logprobs(self, choice_dict: Dict[str, Any]) -> tuple[List[int], List[float]]:
         logprobs_block = choice_dict.get("logprobs")
         if not isinstance(logprobs_block, dict) or not isinstance(logprobs_block.get("content"), list):

responses_api_models/vllm_model/tests/test_app.py

@@ -4679,6 +4679,83 @@ def test_capture_path_can_request_prompt_and_generation_token_ids(self) -> None:
 
         assert result["return_token_ids"] is True
 
+    def test_capture_path_derives_required_prefix_from_latest_assistant(self) -> None:
+        model = _make_top_logprobs_model(return_token_id_information=True)
+        messages = [
+            {"role": "user", "content": "first"},
+            {
+                "role": "assistant",
+                "content": "old",
+                "prompt_token_ids": [1],
+                "generation_token_ids": [2],
+                "generation_log_probs": [-0.1],
+            },
+            {"role": "user", "content": "second"},
+            {
+                "role": "assistant",
+                "content": "new",
+                "prompt_token_ids": [3, 4],
+                "generation_token_ids": [5, 6],
+                "generation_log_probs": [-0.2, -0.3],
+            },
+            {"role": "user", "content": "third"},
+        ]
+
+        result = model._preprocess_chat_completion_create_params(
+            MagicMock(), {"model": "dummy_model", "messages": messages}
+        )
+
+        assert result["required_prefix_token_ids"] == [3, 4, 5, 6]
+
+    def test_capture_path_preserves_explicit_prefix_and_first_turn_omits_it(
+        self,
+    ) -> None:
+        model = _make_top_logprobs_model(return_token_id_information=True)
+        prior_assistant = {
+            "role": "assistant",
+            "content": "answer",
+            "prompt_token_ids": [1],
+            "generation_token_ids": [2],
+            "generation_log_probs": [-0.1],
+        }
+
+        explicit = model._preprocess_chat_completion_create_params(
+            MagicMock(),
+            {
+                "model": "dummy_model",
+                "messages": [prior_assistant],
+                "required_prefix_token_ids": [9, 8],
+            },
+        )
+        first_turn = model._preprocess_chat_completion_create_params(
+            MagicMock(),
+            {
+                "model": "dummy_model",
+                "messages": [{"role": "user", "content": "hi"}],
+            },
+        )
+
+        assert explicit["required_prefix_token_ids"] == [9, 8]
+        assert "required_prefix_token_ids" not in first_turn
+
+    def test_capture_path_rejects_partial_assistant_token_bundle(self) -> None:
+        model = _make_top_logprobs_model(return_token_id_information=True)
+
+        with raises(RuntimeError, match="partial token metadata"):
+            model._preprocess_chat_completion_create_params(
+                MagicMock(),
+                {
+                    "model": "dummy_model",
+                    "messages": [
+                        {
+                            "role": "assistant",
+                            "content": "answer",
+                            "prompt_token_ids": [1],
+                        }
+                    ],
+                },
+            )
+
     def test_capture_path_rejects_multiple_choices(self) -> None:
         model = _make_top_logprobs_model(return_token_id_information=True)