fix(experimental): validate concat parent token prefixes

areal/experimental/openai/client.py

@@ -55,7 +55,10 @@
 from areal.api.cli_args import GenerationHyperparameters
 from areal.experimental.openai.cache import InteractionCache
 from areal.experimental.openai.tool_call_parser import process_tool_calls
-from areal.experimental.openai.types import InteractionWithTokenLogpReward
+from areal.experimental.openai.types import (
+    InteractionWithTokenLogpReward,
+    ensure_parent_token_prefix,
+)
 from areal.utils import logging
 from areal.utils.hf_utils import apply_chat_template
 
@@ -405,6 +408,7 @@ def concat_prompt_token_ids_with_parent(
     Concatenate prompt token IDs with parent interaction's tokens.
     """
     parent_tokens: list[int] = []
+    expected_parent_tokens: list[int] | None = None
     all_message_list: list[dict] = []
     eos_token_id = tokenizer.eos_token_id
 
@@ -421,6 +425,9 @@ def concat_prompt_token_ids_with_parent(
             parent.model_response.input_tokens
             + parent.model_response.output_tokens_without_stop  # without stop tokens
         )
+        expected_parent_tokens = (
+            parent.model_response.input_tokens + parent.model_response.output_tokens
+        )
         all_message_list += parent.messages if parent.messages is not None else []
         all_message_list += (
             parent.output_message_list if parent.output_message_list is not None else []
@@ -462,6 +469,8 @@ def concat_prompt_token_ids_with_parent(
         child_tokens_truncate_idx = -1
 
     prompt_token_ids = parent_tokens + all_tokens[child_tokens_truncate_idx + 1 :]
+    if expected_parent_tokens is not None:
+        ensure_parent_token_prefix(prompt_token_ids, expected_parent_tokens)
     return prompt_token_ids
 
 
@@ -684,13 +693,18 @@ async def create(
                 )
             else:
                 concat_tok_messages = concat_messages
-            prompt_token_ids = concat_prompt_token_ids_with_parent(
-                concat_tok_messages,
-                interaction.parent if interaction is not None else None,
-                self.tokenizer,
-                tools=tools_list,
-                extra_body=extra_body,
-            )
+            try:
+                prompt_token_ids = concat_prompt_token_ids_with_parent(
+                    concat_tok_messages,
+                    interaction.parent if interaction is not None else None,
+                    self.tokenizer,
+                    tools=tools_list,
+                    extra_body=extra_body,
+                )
+            except Exception:
+                if cache is not None:
+                    cache.pop(completion_id, None)
+                raise
         else:
             raise RuntimeError(
                 f"Unsupported chat_template_type {self.chat_template_type}"
@@ -1124,13 +1138,17 @@ async def create(
                 _, remaining_tok, _ = _extract_images_from_messages(remaining)
             else:
                 remaining_tok = remaining
-            prompt_token_ids = concat_prompt_token_ids_with_parent(
-                remaining_tok,
-                interaction.parent if interaction is not None else None,
-                self.tokenizer,
-                tools=tools_list,
-                extra_body=extra_body,
-            )
+            try:
+                prompt_token_ids = concat_prompt_token_ids_with_parent(
+                    remaining_tok,
+                    interaction.parent if interaction is not None else None,
+                    self.tokenizer,
+                    tools=tools_list,
+                    extra_body=extra_body,
+                )
+            except Exception:
+                cache.pop(resp_id, None)
+                raise
         else:
             raise RuntimeError(
                 f"Unsupported chat_template_type {self.chat_template_type}"

areal/experimental/openai/types.py

@@ -16,6 +16,17 @@
 logger = logging.getLogger("TokenLogpReward")
 
 
+def ensure_parent_token_prefix(
+    child_input_tokens: list[int], parent_input_tokens: list[int]
+) -> None:
+    """Require a child prompt to preserve the exact parent token sequence."""
+    if child_input_tokens[: len(parent_input_tokens)] != parent_input_tokens:
+        raise ValueError(
+            "Child input token prefix does not match the parent token "
+            "sequence; refusing to reuse parent token data."
+        )
+
+
 class ApiType(str, Enum):
     """API type for interaction."""
 
@@ -148,12 +159,14 @@ def to_tensor_dict(self) -> dict[str, torch.Tensor]:
         self.seq_tokens = seq = resp.input_tokens + resp.output_tokens
         if self.chat_template_type == "concat" and self.parent is not None:
             parent_res = self.parent.to_tensor_dict()
+            parent_input_ids = parent_res["input_ids"].squeeze(0).tolist()
             parent_logprobs = parent_res["logprobs"].squeeze(0).tolist()
             parent_loss_mask = parent_res["loss_mask"].squeeze(0).tolist()
             parent_versions = parent_res["versions"].squeeze(0).tolist()
             parent_len = len(parent_logprobs)
             assert parent_len == len(parent_loss_mask) == len(parent_versions)
             if resp.input_len > parent_len:
+                ensure_parent_token_prefix(resp.input_tokens, parent_input_ids)
                 logprobs = (
                     parent_logprobs
                     + [0.0] * (resp.input_len - parent_len)

tests/experimental/openai/test_client_cache_lifecycle.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from areal.experimental.openai import ArealOpenAI
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("api", ["chat_completions", "responses"])
+async def test_prompt_construction_failure_discards_cached_interaction(api):
+    """A rejected prompt must not leave an incomplete interaction in the cache."""
+    engine = AsyncMock()
+    client = ArealOpenAI(
+        engine=engine,
+        tokenizer=object(),
+        api_key="test",
+        chat_template_type="concat",
+    )
+
+    with (
+        patch(
+            "areal.experimental.openai.client.concat_prompt_token_ids_with_parent",
+            side_effect=ValueError("invalid parent token prefix"),
+        ),
+        pytest.raises(ValueError, match="invalid parent token prefix"),
+    ):
+        if api == "chat_completions":
+            await client.chat.completions.create(
+                messages=[{"role": "user", "content": "hello"}]
+            )
+        else:
+            await client.responses.create(input="hello")
+
+    assert not client._cache
+    engine.agenerate.assert_not_awaited()

tests/experimental/openai/test_types.py

@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+
+from areal.api import ModelResponse
+from areal.experimental.openai.client import concat_prompt_token_ids_with_parent
+from areal.experimental.openai.types import InteractionWithTokenLogpReward
+
+
+def _make_interaction(
+    input_tokens: list[int],
+    output_tokens: list[int],
+    output_logprobs: list[float],
+    output_versions: list[int],
+    parent: InteractionWithTokenLogpReward | None = None,
+) -> InteractionWithTokenLogpReward:
+    return InteractionWithTokenLogpReward(
+        model_response=ModelResponse(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            output_logprobs=output_logprobs,
+            output_versions=output_versions,
+        ),
+        parent=parent,
+        chat_template_type="concat",
+    )
+
+
+def test_to_tensor_dict_mismatched_parent_prefix_raises():
+    parent = _make_interaction([10], [20], [-0.2], [7])
+    child = _make_interaction([10, 999, 30], [40], [-0.4], [8], parent)
+
+    with pytest.raises(ValueError, match="does not match the parent token sequence"):
+        child.to_tensor_dict()
+
+
+def test_to_tensor_dict_matching_parent_prefix_preserves_parent_data():
+    parent = _make_interaction([10], [20], [-0.2], [7])
+    child = _make_interaction([10, 20, 30], [40], [-0.4], [8], parent)
+
+    result = child.to_tensor_dict()
+
+    expected = {
+        "input_ids": torch.tensor([[10, 20, 30, 40]]),
+        "loss_mask": torch.tensor([[0, 1, 0, 1]]),
+        "logprobs": torch.tensor([[0.0, -0.2, 0.0, -0.4]]),
+        "versions": torch.tensor([[-1, 7, -1, 8]]),
+    }
+    for key, expected_tensor in expected.items():
+        torch.testing.assert_close(result[key], expected_tensor, rtol=0.0, atol=0.0)
+
+
+@pytest.mark.parametrize(
+    ("child_input_tokens", "expected"),
+    [
+        (
+            [10, 20],
+            {
+                "input_ids": torch.tensor([[10, 20, 40]]),
+                "loss_mask": torch.tensor([[0, 0, 1]]),
+                "logprobs": torch.tensor([[0.0, 0.0, -0.4]]),
+                "versions": torch.tensor([[-1, -1, 8]]),
+            },
+        ),
+        (
+            [10],
+            {
+                "input_ids": torch.tensor([[10, 40]]),
+                "loss_mask": torch.tensor([[0, 1]]),
+                "logprobs": torch.tensor([[0.0, -0.4]]),
+                "versions": torch.tensor([[-1, 8]]),
+            },
+        ),
+    ],
+    ids=["equal", "shorter"],
+)
+def test_to_tensor_dict_equal_or_shorter_child_ignores_parent(
+    child_input_tokens: list[int], expected: dict[str, torch.Tensor]
+):
+    parent = _make_interaction([10], [20], [-0.2], [7])
+    child = _make_interaction(child_input_tokens, [40], [-0.4], [8], parent)
+
+    result = child.to_tensor_dict()
+
+    for key, expected_tensor in expected.items():
+        torch.testing.assert_close(result[key], expected_tensor, rtol=0.0, atol=0.0)
+
+
+def test_concat_prompt_rejects_replaced_parent_stop_token(monkeypatch):
+    class _Tokenizer:
+        eos_token_id = 99
+        pad_token_id = 0
+
+    parent = _make_interaction([10], [20, 0], [-0.2, -0.3], [7, 7])
+    parent.model_response.tokenizer = _Tokenizer()
+
+    monkeypatch.setattr(
+        "areal.experimental.openai.client.apply_chat_template",
+        lambda *args, **kwargs: [10, 20, 99, 30],
+    )
+
+    with pytest.raises(ValueError, match="does not match the parent token sequence"):
+        concat_prompt_token_ids_with_parent(
+            message_list=[],
+            parent=parent,
+            tokenizer=_Tokenizer(),
+        )