Deprecate EnvStateMessage in favor of Message.info flag

[jamesbraza]

Summary

EnvStateMessage(Message) carries no fields beyond its parent — the subclass is purely a Python type tag. TaskDatasetServer serializes observations through MessagesAdapter (union of Message | ToolRequestMessage | ToolResponseMessage, no EnvStateMessage), so the subclass identity is erased on the wire. A remote consumer that receives the observation list gets back plain Message instances and cannot tell which one is the env-state snapshot, which makes server-side state-replacement (vs. accumulation) impossible to implement downstream.

This PR moves the marker from class identity into data:

• New contract: Message.info["is_env_state"] = True identifies an env-state observation. The server already ships info via MessagesAdapter.dump_python(obs, context={"include_info": True}), so the flag rides through dump → validate_python without any schema change.
• EnvStateMessage becomes a back-compat shim: its __init__ injects the flag into info and emits a DeprecationWarning. Existing external callers keep working; their observations are now wire-safe.
• NBEnvironment.get_env_state_msg (aviary's only internal caller) is migrated off the subclass to Message.create_message(..., info={"is_env_state": True}). The notebook test's isinstance check is replaced with the info-dict check — same contract a remote consumer would use.
• New tests/test_messages.py::test_env_state_info_flag_roundtrips_through_messages_adapter locks in that the flag survives MessagesAdapter.dump_python → validate_python. A second test asserts the shim warns and that caller-supplied info keys are preserved alongside the injected flag.

Design notes

• Key spelling: "is_env_state" (unnamespaced). Considered "aviary.env_state_observation"; the unnamespaced key was chosen for terseness — collisions with other future Message.info keys are acceptable to manage by code review.
• Deprecation mechanism: runtime warnings.warn(DeprecationWarning, stacklevel=2) only, no typing_extensions.deprecated decorator. Python's default filter dedupes by call site, so the warning fires once per caller location.
• No Message.is_env_state helper: the dict key is the contract; consumers check msg.info and msg.info.get("is_env_state") directly.
• core.py re-export retained: aviary.core.EnvStateMessage still imports, so external code can migrate at its own pace within the deprecation window.
• Alternative considered: a discriminated-union refactor (Literal message_type field + Pydantic Field(discriminator=...)) would preserve class identity end-to-end, but requires every wire-traveling Message subclass to gain a literal field and bumps the JSON schema. The info-flag is non-invasive and round-trips through the existing serialization machinery.

Backwards compatibility

• Aviary clients pre-flag: receivers should treat absence of the key as "not env-state". A one-version window where receivers also accept the subclass keeps server-local code paths working.
• Aviary servers pre-flag: emit observations without the key; consumers fall back to existing behavior.
• Existing JSON payloads: schema-compatible, one new optional key inside the already-optional info dict.

Follow-on (outside this repo)

Downstream consumers that maintain a chat history of observations should filter all entries with info.get("is_env_state") is True except the most recent before re-tokenizing the prompt, replacing the current state-accumulation behavior. That work lives in the consumer's observation-to-dict bridge and history-update step, not here.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR deprecates the EnvStateMessage subclass (which is not preserved through MessagesAdapter serialization) and replaces the env-state “marker” with a wire-safe metadata flag: Message.info["is_env_state"] = True. This allows downstream/remote consumers to reliably identify env-state snapshots after dump/validate round-trips.

Changes:

• Add an is_env_state marker contract via Message.info and make EnvStateMessage a deprecation shim that injects the flag.
• Migrate the notebook environment to emit plain Message instances with the info flag instead of EnvStateMessage.
• Add tests to ensure the info flag survives MessagesAdapter.dump_python(...) -> validate_python(...), and that the shim warns + preserves caller-provided info.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
src/aviary/message.py	Adds ENV_STATE_INFO_KEY and deprecates EnvStateMessage by injecting the env-state info flag on construction.
tests/test_messages.py	Adds unit tests for the deprecation shim and for info-flag round-tripping through MessagesAdapter.
packages/notebook/src/aviary/envs/notebook/env.py	Updates get_env_state_msg() to return Message with info={"is_env_state": True}.
packages/notebook/tests/test_nb_env.py	Updates notebook env test assertions to check the info flag instead of isinstance(..., EnvStateMessage).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.