Deprecate EnvStateMessage in favor of Message.info flag

packages/notebook/src/aviary/envs/notebook/env.py

@@ -11,8 +11,14 @@
 
 import aiodocker
 import nbformat
-from aviary.core import Environment, Message, Messages, Tool, ToolRequestMessage
-from aviary.message import EnvStateMessage
+from aviary.core import (
+    ENV_STATE_INFO_KEY,
+    Environment,
+    Message,
+    Messages,
+    Tool,
+    ToolRequestMessage,
+)
 from jupyter_client.manager import AsyncKernelManager
 from nbformat import NotebookNode
 from numpy.typing import NDArray
@@ -327,20 +333,21 @@ async def _run_notebook_local(self, cell_idx: int | None = None) -> str:
         logger.debug("Reloading notebook from disk")
         return "Executed all cells."
 
-    def get_env_state_msg(self) -> EnvStateMessage:
+    def get_env_state_msg(self) -> Message:
         nb_path = self.state.get_container_path(self.state.nb_path)
         md_notebook, notebook_images = utils.view_notebook(
             cells=self.state.cells, language=self.language.value
         )
         # Write the markdown representation to disk
         self.state.nb_path.with_suffix(".md").write_text(md_notebook)
 
-        return EnvStateMessage.create_message(
+        return Message.create_message(
             text=(
                 "Markdown representation of notebook contents"
                 f" ({nb_path}):\n\n{md_notebook}"
             ),
             images=cast(list[NDArray[Any] | str | bytes], notebook_images),
+            info={ENV_STATE_INFO_KEY: True},
         )
 
     async def close(self):

packages/notebook/tests/test_nb_env.py

@@ -9,7 +9,6 @@
 
 import nbformat.v4 as nbf
 import pytest
-from aviary.message import EnvStateMessage
 
 from aviary.envs.notebook import NBEnvironment
 from aviary.envs.notebook.config import NB_ENVIRONMENT_DOCKER_IMAGE
@@ -121,7 +120,7 @@ async def test_notebook_env(self, use_docker: bool, language: NBLanguage):
             obs, tools = await env.reset()
 
             assert len(obs) == 1
-            assert isinstance(obs[0], EnvStateMessage)
+            assert (obs[0].info or {}).get("is_env_state")
             assert isinstance(obs[0].content, str)
             assert obs[0].content.count("Cell") == 0
 

src/aviary/core.py

@@ -16,7 +16,13 @@
     TaskEnvironmentClient,
 )
 from aviary.functional import DynamicState, fenv
-from aviary.message import EnvStateMessage, MalformedMessageError, Message, join
+from aviary.message import (
+    ENV_STATE_INFO_KEY,
+    EnvStateMessage,
+    MalformedMessageError,
+    Message,
+    join,
+)
 from aviary.render import Renderer
 from aviary.tools import (
     INVALID_TOOL_NAME,
@@ -50,6 +56,7 @@
 
 __all__ = [
     "DEFAULT_EVAL_MODEL_NAME",
+    "ENV_STATE_INFO_KEY",
     "INVALID_TOOL_NAME",
     "TASK_DATASET_REGISTRY",
     "DummyEnv",

src/aviary/message.py

@@ -1,7 +1,8 @@
 import json
 import logging
-from collections.abc import Iterable
-from typing import TYPE_CHECKING, ClassVar, Self
+import warnings
+from collections.abc import Iterable, Mapping, MutableMapping
+from typing import TYPE_CHECKING, Any, ClassVar, Self
 
 from pydantic import (
     BaseModel,
@@ -311,8 +312,34 @@ def common_retryable_errors_log_filter(cls, record: "LogRecord") -> bool:
         return not all(x in record.msg for x in (cls.__name__, EMPTY_CONTENT_BASE_MSG))
 
 
+ENV_STATE_INFO_KEY = "is_env_state"
+
+
 class EnvStateMessage(Message):
-    """A message that contains the current state of the environment."""
+    """A message that contains the current state of the environment.
+
+    Since this subclass is erased by `aviary.tools.base.MessagesAdapter`
+    serialization and it just becomes a Message, we are deprecating this class
+    in favor of a plain Message carrying the `ENV_STATE_INFO_KEY` flag inside `info`,
+    which survives serialization.
+    """
+
+    @model_validator(mode="before")
+    @classmethod
+    def _deprecate_and_inject_env_state_marker(
+        cls, data: MutableMapping[str, Any]
+    ) -> MutableMapping[str, Any]:
+        warnings.warn(
+            f"{cls.__name__} is deprecated; use {Message.__name__}"
+            f"(info={{{ENV_STATE_INFO_KEY!r}: True}}, ...) instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        if isinstance(data, Mapping):
+            info = dict(data.get("info") or {})
+            info[ENV_STATE_INFO_KEY] = True
+            data = {**data, "info": info}
+        return data
 
 
 # Define separately so we can filter out this message type

tests/test_messages.py

@@ -5,7 +5,9 @@
 from lmi import LiteLLMModel
 
 from aviary.core import (
+    EnvStateMessage,
     Message,
+    MessagesAdapter,
     ToolCall,
     ToolCallFunction,
     ToolRequestMessage,
@@ -215,6 +217,28 @@ def test_multimodal_roundtrip_via_string_content(self) -> None:
         assert recovered.content_is_json_str
         assert recovered.is_multimodal
 
+    def test_env_state_message_sets_info_flag_and_warns(self) -> None:
+        with pytest.warns(DeprecationWarning, match="EnvStateMessage"):
+            msg = EnvStateMessage(content="stub")
+        assert (msg.info or {}).get("is_env_state")
+
+        with pytest.warns(DeprecationWarning, match="EnvStateMessage"):
+            msg = EnvStateMessage(content="stub", info={"foo": "bar"})
+        assert (msg.info or {}) == {"foo": "bar", "is_env_state": True}, (
+            "Custom info keys must be preserved alongside the injected flag"
+        )
+
+        with pytest.warns(DeprecationWarning, match="EnvStateMessage"):
+            msg = EnvStateMessage.model_validate({"content": "stub"})
+        assert (msg.info or {}).get("is_env_state")
+
+    def test_env_state_info_flag_roundtrips_through_messages_adapter(self) -> None:
+        original = [Message(content="stub", info={"is_env_state": True})]
+        dumped = MessagesAdapter.dump_python(original, context={"include_info": True})
+        assert dumped[0]["info"]["is_env_state"]
+        (round_trip_original,) = MessagesAdapter.validate_python(dumped)
+        assert (round_trip_original.info or {}).get("is_env_state")
+
     @pytest.mark.parametrize(
         ("images", "message_text", "expected_error", "expected_content_length"),
         [