From 5715a0033e3d2beddae05ea6213bf03c87318923 Mon Sep 17 00:00:00 2001 From: James Braza Date: Thu, 21 May 2026 13:06:32 -0700 Subject: [PATCH 1/2] Deprecate `EnvStateMessage` in favor of `Message.info["is_env_state"]` `EnvStateMessage` carries no data beyond `Message`; its subclass identity is erased by `MessagesAdapter` serialization, so receivers cannot distinguish env-state observations over the wire. Move the marker into `Message.info`, which the server already opts into shipping via `include_info=True`. The subclass keeps working as a shim: its `__init__` injects `info["is_env_state"] = True` and emits a `DeprecationWarning`. A new roundtrip test locks in that the flag survives `MessagesAdapter.dump_python` -> `validate_python`. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/aviary/core.py | 9 ++++++++- src/aviary/message.py | 33 ++++++++++++++++++++++++++++++--- tests/test_messages.py | 24 ++++++++++++++++++++++++ 3 files changed, 62 insertions(+), 4 deletions(-) diff --git a/src/aviary/core.py b/src/aviary/core.py index 8e73bfe0..3a45cdf3 100644 --- a/src/aviary/core.py +++ b/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", diff --git a/src/aviary/message.py b/src/aviary/message.py index c8839722..64fc926d 100644 --- a/src/aviary/message.py +++ b/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 diff --git a/tests/test_messages.py b/tests/test_messages.py index 25e30748..38cb4036 100644 --- a/tests/test_messages.py +++ b/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"), [ From d6efab80d1779b5c79ed54c629eb97683a154418 Mon Sep 17 00:00:00 2001 From: James Braza Date: Thu, 21 May 2026 13:06:43 -0700 Subject: [PATCH 2/2] Migrate `NBEnvironment` off deprecated `EnvStateMessage` `NBEnvironment.get_env_state_msg` now returns a plain `Message` with `info={"is_env_state": True}`, the wire-safe replacement for the subclass tag. The notebook test's `isinstance(..., EnvStateMessage)` check is replaced with the info-dict check, which exercises the same contract a remote consumer would use. Co-Authored-By: Claude Opus 4.7 (1M context) --- packages/notebook/src/aviary/envs/notebook/env.py | 15 +++++++++++---- packages/notebook/tests/test_nb_env.py | 3 +-- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/packages/notebook/src/aviary/envs/notebook/env.py b/packages/notebook/src/aviary/envs/notebook/env.py index 4d81195d..970d1d81 100644 --- a/packages/notebook/src/aviary/envs/notebook/env.py +++ b/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,7 +333,7 @@ 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 @@ -335,12 +341,13 @@ def get_env_state_msg(self) -> EnvStateMessage: # 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): diff --git a/packages/notebook/tests/test_nb_env.py b/packages/notebook/tests/test_nb_env.py index 1635f19f..b5cfe882 100644 --- a/packages/notebook/tests/test_nb_env.py +++ b/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