Replay raw completion prompts without chat wrapping

tests/test_core_flow.py

@@ -5586,6 +5586,42 @@ class CoreFlowTests(unittest.TestCase):
             self.assertEqual(requests[2].body["min_tokens"], 1)
             self.assertEqual(requests[2].body["max_tokens"], 1)
 
+    def test_raw_completion_mode_preserves_trace_prompt_and_endpoint(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            tmp_path = Path(tmp)
+            study_path = _write_study_assets(
+                tmp_path,
+                trace_overrides={"request_mode": "raw_completion"},
+            )
+            trace_path = tmp_path / "trace_windows" / "traces" / "chat_w1.jsonl"
+            rows = [json.loads(line) for line in trace_path.read_text().splitlines()]
+            for idx, row in enumerate(rows):
+                row["prompt"] = f"<|im_start|>user\nraw prompt {idx}<|im_end|>"
+            trace_path.write_text(
+                "".join(json.dumps(row) + "\n" for row in rows),
+                encoding="utf-8",
+            )
+
+            study = load_study_spec(study_path)
+            _, requests = load_trace_requests(study, study_spec_path=study_path)
+
+            self.assertEqual(study.trace.request_mode, "raw_completion")
+            self.assertEqual(requests[0].api_path, "/v1/completions")
+            self.assertEqual(requests[0].body["prompt"], rows[0]["prompt"])
+            self.assertNotIn("messages", requests[0].body)
+
+    def test_raw_completion_mode_requires_prompt(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            study_path = _write_study_assets(
+                Path(tmp),
+                trace_overrides={"request_mode": "raw_completion"},
+            )
+            study = load_study_spec(study_path)
+            with self.assertRaisesRegex(
+                ValueError, "missing prompt required by raw_completion"
+            ):
+                load_trace_requests(study, study_spec_path=study_path)
+
     def test_run_one_request_fails_fixed_length_completion_mismatch(self) -> None:
         request = TraceRequest(
             row_id="r1",
@@ -8815,6 +8851,35 @@ class CoreFlowTests(unittest.TestCase):
         self.assertIsNone(metrics.completion_tokens)
         self.assertEqual(metrics.completion_tokens_source, "none")
 
+    def test_stream_chat_completion_reads_raw_completion_text(self) -> None:
+        class FakeResponse:
+            def __enter__(self):
+                return self
+
+            def __exit__(self, exc_type, exc, traceback):
+                return False
+
+            def __iter__(self):
+                return iter(
+                    [
+                        b'data: {"choices": [{"text": "x"}]}\n',
+                        b'data: {"choices": [], "usage": {"completion_tokens": 1}}\n',
+                        b"data: [DONE]\n",
+                    ]
+                )
+
+        with mock.patch("aituner.http_client._urlopen", return_value=FakeResponse()):
+            metrics = stream_chat_completion(
+                base_url="http://127.0.0.1:8000",
+                body={"model": "m", "prompt": "raw"},
+                timeout_s=1.0,
+                api_path="/v1/completions",
+            )
+
+        self.assertIsNotNone(metrics.ttft_ms)
+        self.assertEqual(metrics.completion_tokens, 1)
+        self.assertEqual(metrics.streamed_chunk_count, 1)
+
     def test_stream_chat_completion_marks_same_instant_multitoken_tpot_unmeasurable(self) -> None:
         class FakeResponse:
             def __enter__(self):