From e4eff34e7835bb4744faf47cb5f55b3d0b171f87 Mon Sep 17 00:00:00 2001
From: The tunix Authors <no-reply@google.com>
Date: Thu, 20 Aug 2026 11:20:32 -0700
Subject: [PATCH] Parse chat messages for string-prompt samplers

PiperOrigin-RevId: 967941432
---
 tests/experimental/rollout/collector_test.py | 50 ++++++++++++++++++++
 tunix/experimental/rollout/collector.py      |  9 +++-
 2 files changed, 58 insertions(+), 1 deletion(-)
 create mode 100644 tests/experimental/rollout/collector_test.py

diff --git a/tests/experimental/rollout/collector_test.py b/tests/experimental/rollout/collector_test.py
new file mode 100644
index 000000000..1fe1856a8
--- /dev/null
+++ b/tests/experimental/rollout/collector_test.py
@@ -0,0 +1,50 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the trajectory collector."""
+
+from absl.testing import absltest
+from tunix.experimental.rollout import collector
+
+
+class _RecordingParser:
+
+  def __init__(self):
+    self.calls = []
+
+  def parse(self, msgs, add_generation_prompt=False, is_first_msg=False):
+    self.calls.append((msgs, add_generation_prompt, is_first_msg))
+    return "PARSED"
+
+
+class BuildPromptTest(absltest.TestCase):
+
+  def test_chat_messages_are_parsed(self):
+    parser = _RecordingParser()
+    msgs = [{"role": "user", "content": "hi"}]
+    self.assertEqual(collector._build_prompt(parser, msgs), "PARSED")
+    self.assertEqual(parser.calls, [(msgs, True, True)])
+
+  def test_string_prompt_passes_through(self):
+    parser = _RecordingParser()
+    self.assertEqual(collector._build_prompt(parser, "raw"), "raw")
+    self.assertEmpty(parser.calls)
+
+  def test_no_parser_passes_through(self):
+    msgs = [{"role": "user", "content": "hi"}]
+    self.assertIs(collector._build_prompt(None, msgs), msgs)
+
+
+if __name__ == "__main__":
+  absltest.main()
\ No newline at end of file
diff --git a/tunix/experimental/rollout/collector.py b/tunix/experimental/rollout/collector.py
index a7da1eba8..5b3917b1f 100644
--- a/tunix/experimental/rollout/collector.py
+++ b/tunix/experimental/rollout/collector.py
@@ -22,6 +22,13 @@
 from tunix.rl.agentic.trajectory import trajectory_collect_engine as rl_collect_engine
 from tunix.rl.rollout import base_rollout
 
+def _build_prompt(chat_parser: Any, chat_completions: Any) -> Any:
+  """Vanilla samplers take a string; parse chat messages when needed."""
+  if chat_parser and not isinstance(chat_completions, str):
+    return chat_parser.parse(
+        chat_completions, add_generation_prompt=True, is_first_msg=True
+    )
+  return chat_completions
 
 class TrajectoryCollectorEngine:
   """Wrapper around TrajectoryCollectEngine providing lifecycle controls and Trajectory conversion."""
@@ -81,7 +88,7 @@ async def model_call(
       )
       sampling_req = sampler_lib.SamplingRequest(
           request_id=self.traj_id,
-          prompt=chat_completions,
+          prompt=_build_prompt(self.chat_parser, chat_completions),
           sampling_params=sampling_params,
       )
       res = await self.sampler.sample(sampling_req, **generation_kwargs)