Loosen Envionment.step signature to accept Message

[sidnarayanan]

Per title, this allows environments to accept non-tool-call Messages. I've added a helper method to cast back to a tool request, so all existing environment behavior is preserved with one extra type check.

[Copilot]

Pull request overview

This PR broadens the step() action type across environments/clients to accept a generic Message, then enforces ToolRequestMessage at runtime via a shared Environment.check_action_is_tool_request() helper. It also expands the server-side /step request schema and adds tests around StepRequest deserialization and non-tool actions.

Changes:

• Update Environment.step() signatures (and multiple concrete envs) to accept Message and validate tool requests at runtime.
• Widen TaskDatasetServer’s StepRequest.action to a union of ToolRequestMessage | ToolResponseMessage | Message.
• Add tests for StepRequest deserialization and for /step calls with non-tool messages.

Reviewed changes

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

Show a summary per file

File	Description
tests/test_envs.py	Adds coverage for StepRequest parsing and /step behavior with non-tool messages
src/aviary/functional.py	Updates functional env wrapper step signature to Message + tool-request validation
src/aviary/env.py	Adds check_action_is_tool_request() and changes abstract step() signature
src/aviary/env_client.py	Allows client step() to accept Message and adjusts error fallback behavior
src/aviary/dataset_server.py	Expands StepRequest schema to accept multiple message types
packages/notebook/src/aviary/envs/notebook/env.py	Updates env step signature + tool-request validation
packages/lfrqa/src/aviary/envs/lfrqa/env.py	Updates env step signature + passes validated tool request to super
packages/labbench/src/aviary/envs/labbench/env.py	Updates env step signature + validates tool request before super call
packages/hotpotqa/src/aviary/envs/hotpotqa/env.py	Updates env step signature + enforces tool-request-only at runtime
packages/gsm8k/src/aviary/envs/gsm8k/env.py	Updates env step signature + enforces tool-request-only at runtime

---

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

[mskarlin]

We're freeing ourselves from the shackles