Merge pull request #159 from kashif/fix-concat

[textarena] flush on newline boundry

Author: burtenshaw

src/envs/textarena_env/server/environment.py

@@ -187,24 +187,24 @@ def _legal_players(self) -> List[int]:
 
     def _convert_messages(self, messages: Iterable[Any]) -> List[TextArenaMessage]:
         converted: List[TextArenaMessage] = []
+        buffered_content: List[str] = []
         buffered_sender: int | None = None
         buffered_category: str | None = None
-        buffered_content: List[str] = []
+        last_char_was_newline = False
 
         def flush_buffer() -> None:
             nonlocal buffered_content, buffered_sender, buffered_category
-            if not buffered_content:
-                return
-            converted.append(
-                TextArenaMessage(
-                    sender_id=buffered_sender if buffered_sender is not None else -1,
-                    content="".join(buffered_content),
-                    category=buffered_category or "MESSAGE",
+            if buffered_content:
+                converted.append(
+                    TextArenaMessage(
+                        sender_id=buffered_sender if buffered_sender is not None else -1,
+                        content="".join(buffered_content),
+                        category=buffered_category or "MESSAGE",
+                    )
                 )
-            )
             buffered_content = []
-            buffered_category = None
             buffered_sender = None
+            buffered_category = None
 
         for entry in messages:
             if isinstance(entry, tuple) and len(entry) == 3:
@@ -219,13 +219,29 @@ def flush_buffer() -> None:
             sender_id = int(sender) if isinstance(sender, (int, float)) else -1
             text = str(content)
 
-            if buffered_content and buffered_category == category_name and buffered_sender == sender_id:
-                buffered_content.append(text)
-            else:
+            if text == "\n":
                 flush_buffer()
+                if last_char_was_newline:
+                    converted.append(
+                        TextArenaMessage(
+                            sender_id=sender_id,
+                            content="",
+                            category=category_name,
+                        )
+                    )
+                last_char_was_newline = True
+                continue
+
+            if buffered_sender is None or buffered_category is None:
                 buffered_sender = sender_id
                 buffered_category = category_name
-                buffered_content = [text]
+            elif buffered_sender != sender_id or buffered_category != category_name:
+                flush_buffer()
+                buffered_sender = sender_id
+                buffered_category = category_name
+
+            buffered_content.append(text)
+            last_char_was_newline = False
 
         flush_buffer()
 

tests/envs/test_textarena_environment.py

@@ -21,3 +21,44 @@ def test_convert_messages_coalesces_consecutive_characters():
         TextArenaMessage(sender_id=1, content="AB", category="MESSAGE"),
         TextArenaMessage(sender_id=2, content="!", category="MESSAGE"),
     ]
+
+
+def test_convert_messages_splits_on_newlines():
+    env = object.__new__(TextArenaEnvironment)
+
+    raw_messages = [
+        "[",
+        "G",
+        "A",
+        "M",
+        "E",
+        "]",
+        "\n",
+        "[",
+        "N",
+        "E",
+        "X",
+        "T",
+        "]",
+    ]
+
+    converted = env._convert_messages(raw_messages)
+
+    assert converted == [
+        TextArenaMessage(sender_id=-1, content="[GAME]", category="MESSAGE"),
+        TextArenaMessage(sender_id=-1, content="[NEXT]", category="MESSAGE"),
+    ]
+
+
+def test_convert_messages_preserves_blank_lines():
+    env = object.__new__(TextArenaEnvironment)
+
+    raw_messages = ["A", "\n", "\n", "B"]
+
+    converted = env._convert_messages(raw_messages)
+
+    assert converted == [
+        TextArenaMessage(sender_id=-1, content="A", category="MESSAGE"),
+        TextArenaMessage(sender_id=-1, content="", category="MESSAGE"),
+        TextArenaMessage(sender_id=-1, content="B", category="MESSAGE"),
+    ]