Fix interception proxy capturing 0 token ids on current vLLM

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]