add env param to command type

packages/runtimeuse-client-python/README.md

@@ -114,6 +114,7 @@ result = await client.query(
     options=QueryOptions(
         system_prompt="You are a helpful assistant.",
         model="gpt-4.1",
+        agent_env={"MY_VAR": "value"},               # optional -- env vars for the agent
         pre_agent_downloadables=[downloadable],          # optional
         output_format_json_schema_str='...',         # optional -- omit for text response
         on_assistant_message=on_assistant,            # optional
@@ -147,7 +148,7 @@ result = await client.execute_commands(
     commands=[
         CommandInterface(command="mkdir -p /app/output"),
         CommandInterface(command="echo 'sandbox is ready' > /app/output/status.txt"),
-        CommandInterface(command="cat /app/output/status.txt"),
+        CommandInterface(command="cat /app/output/status.txt", env={"MY_VAR": "value"}),
     ],
     options=ExecuteCommandsOptions(
         on_assistant_message=on_assistant,  # optional -- streams stdout/stderr
@@ -203,22 +204,22 @@ except CancelledException:
 
 ### Types
 
-| Class                                     | Description                                                              |
-| ----------------------------------------- | ------------------------------------------------------------------------ |
-| `QueryOptions`                            | Configuration for `client.query()` (prompt options, callbacks, timeout)  |
-| `QueryResult`                             | Return type of `query()` (`.data`, `.metadata`)                          |
-| `ResultMessageInterface`                  | Wire-format result message from the runtime                              |
-| `TextResult`                              | Result variant when no output schema is specified (`.text`)              |
-| `StructuredOutputResult`                  | Result variant when an output schema is specified (`.structured_output`) |
-| `AssistantMessageInterface`               | Intermediate assistant text messages                                     |
-| `ArtifactUploadRequestMessageInterface`   | Runtime requesting a presigned URL for artifact upload                   |
-| `ArtifactUploadResponseMessageInterface`  | Response with presigned URL sent back to runtime                         |
-| `ErrorMessageInterface`                   | Error from the agent runtime                                             |
-| `ExecuteCommandsOptions`                  | Configuration for `client.execute_commands()` (callbacks, timeout)        |
-| `CommandExecutionResult`                  | Return type of `execute_commands()` (`.results`)                         |
-| `CommandResultItem`                       | Per-command result (`.command`, `.exit_code`)                            |
-| `CommandInterface`                        | Shell command to execute (`.command`, `.cwd`)                            |
-| `RuntimeEnvironmentDownloadableInterface` | File to download into the runtime before invocation                      |
+| Class                                     | Description                                                                          |
+| ----------------------------------------- | ------------------------------------------------------------------------------------ |
+| `QueryOptions`                            | Configuration for `client.query()` (prompt options, `agent_env`, callbacks, timeout) |
+| `QueryResult`                             | Return type of `query()` (`.data`, `.metadata`)                                      |
+| `ResultMessageInterface`                  | Wire-format result message from the runtime                                          |
+| `TextResult`                              | Result variant when no output schema is specified (`.text`)                          |
+| `StructuredOutputResult`                  | Result variant when an output schema is specified (`.structured_output`)             |
+| `AssistantMessageInterface`               | Intermediate assistant text messages                                                 |
+| `ArtifactUploadRequestMessageInterface`   | Runtime requesting a presigned URL for artifact upload                               |
+| `ArtifactUploadResponseMessageInterface`  | Response with presigned URL sent back to runtime                                     |
+| `ErrorMessageInterface`                   | Error from the agent runtime                                                         |
+| `ExecuteCommandsOptions`                  | Configuration for `client.execute_commands()` (callbacks, timeout)                   |
+| `CommandExecutionResult`                  | Return type of `execute_commands()` (`.results`)                                     |
+| `CommandResultItem`                       | Per-command result (`.command`, `.exit_code`)                                        |
+| `CommandInterface`                        | Shell command to execute (`.command`, `.cwd`, `.env`)                                |
+| `RuntimeEnvironmentDownloadableInterface` | File to download into the runtime before invocation                                  |
 
 ### Exceptions
 

packages/runtimeuse-client-python/src/runtimeuse_client/client.py

@@ -98,6 +98,7 @@ async def query(
             model=options.model,
             output_format_json_schema_str=options.output_format_json_schema_str,
             source_id=options.source_id,
+            agent_env=options.agent_env,
             secrets_to_redact=options.secrets_to_redact,
             artifacts_dir=options.artifacts_dir,
             pre_agent_invocation_commands=options.pre_agent_invocation_commands,
@@ -256,21 +257,16 @@ async def execute_commands(
                     raise CancelledException("Command execution was cancelled")
 
                 try:
-                    message_interface = AgentRuntimeMessageInterface.model_validate(
-                        msg
-                    )
+                    message_interface = AgentRuntimeMessageInterface.model_validate(msg)
                 except pydantic.ValidationError:
                     logger.error(
                         f"Received unknown message type from agent runtime: {msg}"
                     )
                     continue
 
-                if (
-                    message_interface.message_type
-                    == "command_execution_result_message"
-                ):
-                    wire_result = (
-                        CommandExecutionResultMessageInterface.model_validate(msg)
+                if message_interface.message_type == "command_execution_result_message":
+                    wire_result = CommandExecutionResultMessageInterface.model_validate(
+                        msg
                     )
                     logger.info(
                         f"Received command execution result from agent runtime: {msg}"
@@ -282,15 +278,13 @@ async def execute_commands(
                         assistant_message_interface = (
                             AssistantMessageInterface.model_validate(msg)
                         )
-                        await options.on_assistant_message(
-                            assistant_message_interface
-                        )
+                        await options.on_assistant_message(assistant_message_interface)
                     continue
 
                 elif message_interface.message_type == "error_message":
                     try:
-                        error_message_interface = (
-                            ErrorMessageInterface.model_validate(msg)
+                        error_message_interface = ErrorMessageInterface.model_validate(
+                            msg
                         )
                     except pydantic.ValidationError:
                         logger.error(
@@ -306,17 +300,14 @@ async def execute_commands(
                     )
 
                 elif (
-                    message_interface.message_type
-                    == "artifact_upload_request_message"
+                    message_interface.message_type == "artifact_upload_request_message"
                 ):
                     logger.info(
                         f"Received artifact upload request message from agent runtime: {msg}"
                     )
                     if options.on_artifact_upload_request is not None:
                         artifact_upload_request_message_interface = (
-                            ArtifactUploadRequestMessageInterface.model_validate(
-                                msg
-                            )
+                            ArtifactUploadRequestMessageInterface.model_validate(msg)
                         )
                         upload_result = await options.on_artifact_upload_request(
                             artifact_upload_request_message_interface

packages/runtimeuse-client-python/src/runtimeuse_client/types.py

@@ -26,13 +26,15 @@ class CommandInterface(BaseModel):
 
     cwd: str | None = None
     command: str
+    env: dict[str, str] | None = None
 
 
 class InvocationMessage(BaseModel):
     message_type: Literal["invocation_message"]
     source_id: str | None = None
     system_prompt: str
     user_prompt: str
+    agent_env: dict[str, str] | None = None
     output_format_json_schema_str: str | None = None
     secrets_to_redact: list[str] = Field(default_factory=list)
     artifacts_dir: str | None = None
@@ -106,7 +108,9 @@ class CommandExecutionMessage(BaseModel):
     secrets_to_redact: list[str] = Field(default_factory=list)
     commands: list[CommandInterface]
     artifacts_dir: str | None = None
-    pre_execution_downloadables: list[RuntimeEnvironmentDownloadableInterface] | None = None
+    pre_execution_downloadables: (
+        list[RuntimeEnvironmentDownloadableInterface] | None
+    ) = None
 
 
 class CommandResultItem(BaseModel):
@@ -157,6 +161,8 @@ class QueryOptions:
 
     #: Caller-defined identifier for tracing/logging purposes.
     source_id: str | None = None
+    #: Environment variables to set in the agent runtime.
+    agent_env: dict[str, str] | None = None
     #: Secret values to redact from agent logs and responses.
     secrets_to_redact: list[str] = field(default_factory=list)
     #: Directory inside the runtime environment where artifacts are written.
@@ -197,7 +203,9 @@ class ExecuteCommandsOptions:
     #: Directory inside the runtime environment where artifacts are written.
     artifacts_dir: str | None = None
     #: Files to download into the runtime environment before commands run.
-    pre_execution_downloadables: list[RuntimeEnvironmentDownloadableInterface] | None = None
+    pre_execution_downloadables: (
+        list[RuntimeEnvironmentDownloadableInterface] | None
+    ) = None
     #: Called for each assistant (intermediate) message streamed back.
     on_assistant_message: OnAssistantMessageCallback | None = None
     #: Called when the runtime requests an artifact upload URL.

packages/runtimeuse-client-python/test/test_client.py

@@ -45,7 +45,10 @@ class TestResultMessage:
     async def test_structured_output_result(self, fake_transport, make_query_options):
         result_msg = {
             "message_type": "result_message",
-            "data": {"type": "structured_output", "structured_output": {"success": True}},
+            "data": {
+                "type": "structured_output",
+                "structured_output": {"success": True},
+            },
             "metadata": {"duration_ms": 50},
         }
         transport, client = fake_transport([result_msg])
@@ -363,7 +366,10 @@ async def test_full_message_sequence(self, fake_transport, make_query_options):
             },
             {
                 "message_type": "result_message",
-                "data": {"type": "structured_output", "structured_output": {"answer": 42}},
+                "data": {
+                    "type": "structured_output",
+                    "structured_output": {"answer": 42},
+                },
                 "metadata": {"duration_ms": 100},
             },
         ]
@@ -421,6 +427,36 @@ async def test_schema_forwarded_when_set(self, fake_transport, make_query_option
             invocation_msgs[0]["output_format_json_schema_str"] == '{"type":"object"}'
         )
 
+    @pytest.mark.asyncio
+    async def test_agent_env_forwarded_when_set(
+        self, fake_transport, make_query_options
+    ):
+        transport, client = fake_transport([TEXT_RESULT_MSG])
+
+        await client.query(
+            prompt=DEFAULT_PROMPT,
+            options=make_query_options(agent_env={"MY_VAR": "hello"}),
+        )
+
+        invocation_msgs = [
+            m for m in transport.sent if m.get("message_type") == "invocation_message"
+        ]
+        assert invocation_msgs[0]["agent_env"] == {"MY_VAR": "hello"}
+
+    @pytest.mark.asyncio
+    async def test_agent_env_none_when_omitted(self, fake_transport, query_options):
+        transport, client = fake_transport([TEXT_RESULT_MSG])
+
+        await client.query(
+            prompt=DEFAULT_PROMPT,
+            options=query_options,
+        )
+
+        invocation_msgs = [
+            m for m in transport.sent if m.get("message_type") == "invocation_message"
+        ]
+        assert invocation_msgs[0]["agent_env"] is None
+
     @pytest.mark.asyncio
     async def test_schema_none_when_omitted(self, fake_transport, query_options):
         transport, client = fake_transport([TEXT_RESULT_MSG])
@@ -518,7 +554,9 @@ async def test_sends_command_execution_message(
         ]
         assert len(cmd_msgs) == 1
         assert cmd_msgs[0]["source_id"] == "cmd-test"
-        assert cmd_msgs[0]["commands"] == [{"command": "echo hello", "cwd": None}]
+        assert cmd_msgs[0]["commands"] == [
+            {"command": "echo hello", "cwd": None, "env": None}
+        ]
 
     @pytest.mark.asyncio
     async def test_assistant_message_dispatched(
@@ -641,6 +679,49 @@ async def on_artifact(
         assert response_msgs[0]["filename"] == "output.txt"
         assert response_msgs[0]["presigned_url"] == "https://s3.example.com/presigned"
 
+    @pytest.mark.asyncio
+    async def test_command_env_forwarded(
+        self, fake_transport, make_execute_commands_options
+    ):
+        result_msg = {
+            "message_type": "command_execution_result_message",
+            "results": [{"command": "echo hello", "exit_code": 0}],
+        }
+        transport, client = fake_transport([result_msg])
+
+        await client.execute_commands(
+            commands=[CommandInterface(command="echo hello", env={"FOO": "bar"})],
+            options=make_execute_commands_options(),
+        )
+
+        cmd_msgs = [
+            m
+            for m in transport.sent
+            if m.get("message_type") == "command_execution_message"
+        ]
+        assert len(cmd_msgs) == 1
+        assert cmd_msgs[0]["commands"] == [
+            {"command": "echo hello", "cwd": None, "env": {"FOO": "bar"}}
+        ]
+
+    @pytest.mark.asyncio
+    async def test_command_env_none_by_default(
+        self, fake_transport, make_execute_commands_options
+    ):
+        transport, client = fake_transport([COMMAND_RESULT_MSG])
+
+        await client.execute_commands(
+            commands=[CommandInterface(command="echo hello")],
+            options=make_execute_commands_options(),
+        )
+
+        cmd_msgs = [
+            m
+            for m in transport.sent
+            if m.get("message_type") == "command_execution_message"
+        ]
+        assert cmd_msgs[0]["commands"][0]["env"] is None
+
     def test_execute_commands_options_artifacts_validation(self):
         with pytest.raises(ValueError, match="must be specified together"):
             ExecuteCommandsOptions(artifacts_dir="/tmp/artifacts")