Fix interception proxy capturing 0 token ids on current vLLM

vLLM with --return-tokens-as-token-ids encodes the token id in the logprobs
"token" field as "token_id:<int>" rather than a separate "token_id" key,
so the capture read zero token ids. Every env that trains on the captured trace
(opencode, pi, ...) then loses the token ids GRPO needs. Parse the id from the
token field. Adds a unit test.

Author: sergiopaniego

envs/opencode_env/sandbox/interception.py

@@ -130,9 +130,7 @@ async def chat_completions(request: Request) -> Response:
         try:
             body = json.loads(raw_body)
         except json.JSONDecodeError:
-            return JSONResponse(
-                status_code=400, content={"error": "invalid json body"}
-            )
+            return JSONResponse(status_code=400, content={"error": "invalid json body"})
 
         forwarded_body = _prepare_forwarded_body(body, cfg)
         headers = {
@@ -338,7 +336,7 @@ async def _stream() -> Any:
                 yield line + "\n"
                 if not line.startswith("data:"):
                     continue
-                data = line[len("data:"):].strip()
+                data = line[len("data:") :].strip()
                 if data == "[DONE]":
                     continue
                 try:
@@ -381,7 +379,11 @@ def _accumulate_stream_chunk(chunk: dict[str, Any], acc: dict[str, Any]) -> None
             tc_idx = tc.get("index", 0)
             bucket = acc["tool_calls_by_idx"].setdefault(
                 (idx, tc_idx),
-                {"id": None, "type": "function", "function": {"name": "", "arguments": ""}},
+                {
+                    "id": None,
+                    "type": "function",
+                    "function": {"name": "", "arguments": ""},
+                },
             )
             if tc.get("id"):
                 bucket["id"] = tc["id"]
@@ -458,11 +460,11 @@ def _build_turn_record(
     token_ids: list[int] = []
     per_token_logps: list[float] = []
     for entry in content_lp:
-        tokens.append(entry.get("token", ""))
-        # OpenAI returns no raw token ids; vLLM returns them as ``token_id``.
-        token_id = entry.get("token_id")
-        if token_id is not None:
-            token_ids.append(int(token_id))
+        token = entry.get("token", "")
+        tokens.append(token)
+        # vLLM (--return-tokens-as-token-ids) returns the id in the token field as "token_id:<int>".
+        if token.startswith("token_id:"):
+            token_ids.append(int(token.split(":", 1)[1]))
         lp = entry.get("logprob")
         if lp is not None:
             per_token_logps.append(float(lp))
@@ -487,8 +489,7 @@ def _strip_logprobs(response_json: dict[str, Any]) -> dict[str, Any]:
     choices = out.get("choices")
     if isinstance(choices, list):
         out["choices"] = [
-            {k: v for k, v in (ch or {}).items() if k != "logprobs"}
-            for ch in choices
+            {k: v for k, v in (ch or {}).items() if k != "logprobs"} for ch in choices
         ]
     return out
 
@@ -537,9 +538,7 @@ def start(self) -> None:
             lifespan="on",
         )
         self._server = uvicorn.Server(config)
-        self._thread = threading.Thread(
-            target=self._run_server, daemon=True
-        )
+        self._thread = threading.Thread(target=self._run_server, daemon=True)
         self._thread.start()
         # Wait for the server to accept connections.
         deadline = time.time() + 10

tests/envs/test_interception_capture.py

@@ -0,0 +1,51 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""The interception proxy must capture per-token ids from a vLLM logprobs response.
+
+vLLM with ``--return-tokens-as-token-ids`` encodes the id in the ``token`` field as
+``"token_id:<int>"`` (no separate ``token_id`` key), which the capture must parse so
+GRPO trains on real token ids."""
+
+import os
+import sys
+
+# Make ``envs/`` importable when running from the repository root.
+_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+if _REPO_ROOT not in sys.path:
+    sys.path.insert(0, _REPO_ROOT)
+_ENVS_DIR = os.path.join(_REPO_ROOT, "envs")
+if _ENVS_DIR not in sys.path:
+    sys.path.insert(0, _ENVS_DIR)
+
+from opencode_env.sandbox.interception import _build_turn_record
+
+
+def _response(content):
+    return {"choices": [{"finish_reason": "stop", "logprobs": {"content": content}}]}
+
+
+def test_captures_ids_from_vllm_token_id_prefix():
+    # vLLM --return-tokens-as-token-ids: id lives in the ``token`` field as "token_id:N".
+    content = [
+        {"token": "token_id:15339", "logprob": -0.1},
+        {"token": "token_id:1917", "logprob": -0.2},
+    ]
+    rec = _build_turn_record(
+        turn_idx=0, request_body={}, response_json=_response(content), latency_s=0.0
+    )
+    assert rec.completion_token_ids == [15339, 1917]
+    assert rec.per_token_logps == [-0.1, -0.2]
+
+
+def test_plain_openai_has_logprobs_but_no_ids():
+    # A plain OpenAI response (no ids) still yields logprobs, just no token ids.
+    content = [{"token": "hi", "logprob": -0.3}]
+    rec = _build_turn_record(
+        turn_idx=0, request_body={}, response_json=_response(content), latency_s=0.0
+    )
+    assert rec.completion_token_ids == []
+    assert rec.per_token_logps == [-0.3]