diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index 27dd94615dd..8465cb9728e 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -716,6 +716,18 @@ async def chat_completions(request: Request, session_id: str): completion_tokens_emit = -1 _stats["turns_completed"] += 1 + # Strip routed_experts/token-id blobs (in meta_info + prompt_token_ids) + # from the agent's copy only; the trainer reads them from the recorded + # SessionRecord. Returning them every turn made harbor RSS grow unbounded. + result["response_body"] = orjson.dumps( + { + **response, + "choices": [ + {key: value for key, value in choice.items() if key not in ("meta_info", "prompt_token_ids")} + for choice in response.get("choices", []) + ], + } + ) return backend.build_proxy_response(result) finally: _inflight_chat["count"] -= 1 diff --git a/tests/fast/router/test_session_pretokenized_e2e.py b/tests/fast/router/test_session_pretokenized_e2e.py index 68917a68f17..8e676a9bb27 100644 --- a/tests/fast/router/test_session_pretokenized_e2e.py +++ b/tests/fast/router/test_session_pretokenized_e2e.py @@ -9,8 +9,8 @@ prefix of full_str) and uses pretokenized_token_ids for prompt construction. - We check mock server request_log to verify pretokenized fields are injected on turn 2+. -- We verify pretokenized_token_ids is a valid prefix of the returned - prompt_token_ids. +- We verify pretokenized_token_ids is a valid prefix of the prompt_token_ids + retained in the session record. Tests are parametrized across multiple models and chat template combinations. Models whose native templates satisfy the prefix invariant (Qwen3.5, @@ -239,13 +239,14 @@ def _run_trajectory_e2e(env): return session_id, responses -def _verify_pretokenized_injection(env, responses): +def _verify_pretokenized_injection(env, session_id): """Verify input_ids in request_log and prompt consistency.""" backend = env.backend + records = requests.get(f"{env.url}/sessions/{session_id}", timeout=5.0).json()["records"] for turn_idx in range(len(env.trajectory.turns)): req = backend.request_log[turn_idx] - resp_prompt_ids = responses[turn_idx]["choices"][0]["prompt_token_ids"] + resp_prompt_ids = records[turn_idx]["response"]["choices"][0]["prompt_token_ids"] if turn_idx == 0: assert "input_ids" not in req, "Turn 0 should not have input_ids" @@ -262,13 +263,14 @@ def _verify_pretokenized_injection(env, responses): ) -def _verify_pretokenized_NOT_prefix(env, responses): +def _verify_pretokenized_NOT_prefix(env, session_id): """Verify that input_ids does NOT match prompt on turn 2+. This is the inverse of _verify_pretokenized_injection: it proves that a broken native template causes token mismatch in the e2e flow. """ backend = env.backend + records = requests.get(f"{env.url}/sessions/{session_id}", timeout=5.0).json()["records"] found_mismatch = False for turn_idx in range(1, len(env.trajectory.turns)): @@ -276,7 +278,7 @@ def _verify_pretokenized_NOT_prefix(env, responses): if "input_ids" not in req: continue input_ids = req["input_ids"] - resp_prompt_ids = responses[turn_idx]["choices"][0]["prompt_token_ids"] + resp_prompt_ids = records[turn_idx]["response"]["choices"][0]["prompt_token_ids"] if resp_prompt_ids != input_ids: found_mismatch = True @@ -308,8 +310,8 @@ def test_all_turns_succeed(self): assert len(responses) == 2 def test_pretokenized_injection(self): - session_id, responses = _run_trajectory_e2e(self._env) - _verify_pretokenized_injection(self._env, responses) + session_id, _responses = _run_trajectory_e2e(self._env) + _verify_pretokenized_injection(self._env, session_id) def test_session_records(self): session_id, responses = _run_trajectory_e2e(self._env) @@ -336,8 +338,8 @@ def test_all_turns_succeed(self): def test_pretokenized_injection(self): self._env.backend.reset_stats() - session_id, responses = _run_trajectory_e2e(self._env) - _verify_pretokenized_injection(self._env, responses) + session_id, _responses = _run_trajectory_e2e(self._env) + _verify_pretokenized_injection(self._env, session_id) def test_pretokenized_grows_across_turns(self): self._env.backend.reset_stats() diff --git a/tests/fast/router/test_sessions.py b/tests/fast/router/test_sessions.py index 875f5116830..fa4dd8fe4c2 100644 --- a/tests/fast/router/test_sessions.py +++ b/tests/fast/router/test_sessions.py @@ -141,6 +141,9 @@ def test_proxy_chat_appends_record(self, router_env): body = resp.json() assert "choices" in body assert body["choices"] + agent_choice = body["choices"][0] + assert "meta_info" not in agent_choice + assert "prompt_token_ids" not in agent_choice get_resp = requests.get(f"{router_env.url}/sessions/{session_id}", timeout=5.0) records = get_resp.json()["records"] @@ -150,6 +153,16 @@ def test_proxy_chat_appends_record(self, router_env): record = records[0] assert record["path"] == "/v1/chat/completions" assert record["status_code"] == 200 + recorded_choice = record["response"]["choices"][0] + assert recorded_choice["prompt_token_ids"] + assert recorded_choice["meta_info"]["output_token_logprobs"] + + merged_resp = requests.get(f"{router_env.url}/sessions/{session_id}/merged", timeout=5.0) + assert merged_resp.status_code == 200 + sample = merged_resp.json()["sample"] + assert sample is not None + assert sample["response_length"] > 0 + assert len(sample["rollout_log_probs"]) == sample["response_length"] class TestTokenizationOffload: