diff --git a/packages/runtimeuse-client-python/README.md b/packages/runtimeuse-client-python/README.md index 7eaeb2e..74b111d 100644 --- a/packages/runtimeuse-client-python/README.md +++ b/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 diff --git a/packages/runtimeuse-client-python/src/runtimeuse_client/client.py b/packages/runtimeuse-client-python/src/runtimeuse_client/client.py index 992f4ea..032129e 100644 --- a/packages/runtimeuse-client-python/src/runtimeuse_client/client.py +++ b/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 diff --git a/packages/runtimeuse-client-python/src/runtimeuse_client/types.py b/packages/runtimeuse-client-python/src/runtimeuse_client/types.py index 0eec5f5..55704f7 100644 --- a/packages/runtimeuse-client-python/src/runtimeuse_client/types.py +++ b/packages/runtimeuse-client-python/src/runtimeuse_client/types.py @@ -26,6 +26,7 @@ class CommandInterface(BaseModel): cwd: str | None = None command: str + env: dict[str, str] | None = None class InvocationMessage(BaseModel): @@ -33,6 +34,7 @@ class InvocationMessage(BaseModel): 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. diff --git a/packages/runtimeuse-client-python/test/test_client.py b/packages/runtimeuse-client-python/test/test_client.py index 445d7d9..90a340a 100644 --- a/packages/runtimeuse-client-python/test/test_client.py +++ b/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") diff --git a/packages/runtimeuse/README.md b/packages/runtimeuse/README.md index 79adc5e..c632cad 100644 --- a/packages/runtimeuse/README.md +++ b/packages/runtimeuse/README.md @@ -91,15 +91,16 @@ interface AgentHandler { **`AgentInvocation`** -- everything your agent needs: -| Field | Type | Description | -| -------------- | ---------------------------------------------------------- | ------------------------------------ | -| `systemPrompt` | `string` | System prompt for the agent | -| `userPrompt` | `string` | User prompt / task description | -| `outputFormat` | `{ type: "json_schema"; schema: Record }` | Expected output schema | -| `model` | `string` | Model identifier | -| `secrets` | `string[]` | Values to redact from logs | -| `signal` | `AbortSignal` | Observe for cancellation (read-only) | -| `logger` | `Logger` | Prefixed logger for this invocation | +| Field | Type | Description | +| -------------- | ---------------------------------------------------------- | ------------------------------------------ | +| `systemPrompt` | `string` | System prompt for the agent | +| `userPrompt` | `string` | User prompt / task description | +| `outputFormat` | `{ type: "json_schema"; schema: Record }` | Expected output schema | +| `model` | `string` | Model identifier | +| `env` | `Record` (optional) | Environment variables to pass to the agent | +| `secrets` | `string[]` | Values to redact from logs | +| `signal` | `AbortSignal` | Observe for cancellation (read-only) | +| `logger` | `Logger` | Prefixed logger for this invocation | **`MessageSender`** -- send intermediate messages back to the client: @@ -170,14 +171,21 @@ wss.on("connection", (ws) => { When a client sends an `invocation_message`, the session: 1. **Downloads runtime files** -- if `pre_agent_downloadables` is set, fetches and extracts them -2. **Runs pre-commands** -- if `pre_agent_invocation_commands` is set, executes them. If it exits 0, execution continues to the next command or the agent. Any other non-zero exit code sends an error message and terminates the invocation. -3. **Calls `handler.run()`** -- your agent logic runs with the invocation context and a `MessageSender` +2. **Runs pre-commands** -- if `pre_agent_invocation_commands` is set, executes them. Each command can specify its own `env` and `cwd`. If it exits 0, execution continues to the next command or the agent. Any other non-zero exit code sends an error message and terminates the invocation. +3. **Calls `handler.run()`** -- your agent logic runs with the invocation context (including any `agent_env` environment variables) and a `MessageSender` 4. **Sends `result_message`** -- the `AgentResult` from your handler is sent back to the client 5. **Finalizes** -- stops artifact watching, waits for pending uploads, closes the WebSocket -### Command-Only Execution +### Command-Only Execution (no agent) -The session also accepts a `command_execution_message` instead of an `invocation_message`. This runs `pre_execution_downloadables` and the provided commands, streams output as `assistant_message`s, and returns a `command_execution_result_message` with per-command exit codes -- without invoking the agent handler. See the [Python client docs](../runtimeuse-client-python/README.md#command-only-execution) for usage. +The session also accepts a `command_execution_message` instead of an `invocation_message`. This runs `pre_execution_downloadables` and the provided commands, streams output as `assistant_message`s, and returns a `command_execution_result_message` with per-command exit codes. The agent handler never gets invoked. Each command can specify its own `env` and `cwd`. See the [Python client docs](../runtimeuse-client-python/README.md#command-only-execution) for usage. + +## Environment Variables + +Environment variables can be injected at two levels: + +- **Per-command (`Command.env`)** -- each command in `pre_agent_invocation_commands`, `post_agent_invocation_commands`, or `command_execution_message.commands` can carry its own `env` map. These are merged on top of `process.env` when the command is spawned. +- **Per-invocation (`InvocationMessage.agent_env`)** -- environment variables passed to the agent handler. The Claude handler merges these on top of `process.env` when calling the Claude Agent SDK. Custom handlers receive these via `AgentInvocation.env`. ## Artifact Management @@ -224,17 +232,17 @@ Command output (stdout/stderr) from pre-commands is automatically redacted using ### Protocol Message Types -| Type | Direction | Description | -| ------------------------------- | ----------------- | --------------------------------------- | -| `InvocationMessage` | Client -> Runtime | Start an agent invocation | -| `CommandExecutionMessage` | Client -> Runtime | Run commands without agent invocation | -| `CancelMessage` | Client -> Runtime | Cancel a running invocation or execution | -| `ArtifactUploadResponseMessage` | Client -> Runtime | Presigned URL for artifact upload | -| `ResultMessage` | Runtime -> Client | Structured agent result | -| `CommandExecutionResultMessage` | Runtime -> Client | Per-command exit codes | -| `AssistantMessage` | Runtime -> Client | Intermediate text from the agent | -| `ArtifactUploadRequestMessage` | Runtime -> Client | Request a presigned URL for an artifact | -| `ErrorMessage` | Runtime -> Client | Error during execution | +| Type | Direction | Description | +| ------------------------------- | ----------------- | ---------------------------------------- | +| `InvocationMessage` | Client -> Runtime | Start an agent invocation | +| `CommandExecutionMessage` | Client -> Runtime | Run commands without agent invocation | +| `CancelMessage` | Client -> Runtime | Cancel a running invocation or execution | +| `ArtifactUploadResponseMessage` | Client -> Runtime | Presigned URL for artifact upload | +| `ResultMessage` | Runtime -> Client | Structured agent result | +| `CommandExecutionResultMessage` | Runtime -> Client | Per-command exit codes | +| `AssistantMessage` | Runtime -> Client | Intermediate text from the agent | +| `ArtifactUploadRequestMessage` | Runtime -> Client | Request a presigned URL for an artifact | +| `ErrorMessage` | Runtime -> Client | Error during execution | ## Related Docs diff --git a/packages/runtimeuse/src/agent-handler.ts b/packages/runtimeuse/src/agent-handler.ts index 9245735..138e590 100644 --- a/packages/runtimeuse/src/agent-handler.ts +++ b/packages/runtimeuse/src/agent-handler.ts @@ -8,11 +8,16 @@ export interface AgentInvocation { secrets: string[]; signal: AbortSignal; logger: Logger; + env?: Record; } export type AgentResult = | { type: "text"; text: string; metadata?: Record } - | { type: "structured_output"; structuredOutput: Record; metadata?: Record }; + | { + type: "structured_output"; + structuredOutput: Record; + metadata?: Record; + }; export interface MessageSender { sendAssistantMessage(textBlocks: string[]): void; diff --git a/packages/runtimeuse/src/claude-handler.ts b/packages/runtimeuse/src/claude-handler.ts index 0a6ff91..73ef621 100644 --- a/packages/runtimeuse/src/claude-handler.ts +++ b/packages/runtimeuse/src/claude-handler.ts @@ -1,4 +1,4 @@ -import { query } from "@anthropic-ai/claude-agent-sdk"; +import { Options, query } from "@anthropic-ai/claude-agent-sdk"; import type { AgentHandler, AgentInvocation, @@ -36,27 +36,14 @@ export const claudeHandler: AgentHandler = { invocation.signal.addEventListener("abort", onAbort, { once: true }); try { - const queryOptions: Record = { + const queryOptions: Options = { systemPrompt: invocation.systemPrompt, model: invocation.model, abortController, tools: { type: "preset", preset: "claude_code" }, permissionMode: "bypassPermissions", allowDangerouslySkipPermissions: true, - hooks: { - PostToolUse: [ - { - hooks: [ - async (input: Record) => { - invocation.logger.log( - `[PostToolUse] tool=${input.tool_name} input=${JSON.stringify(input.tool_input)} response=${JSON.stringify(input.tool_response)}`, - ); - return {}; - }, - ], - }, - ], - }, + env: { ...process.env, ...invocation.env }, }; if (invocation.outputFormat) { queryOptions.outputFormat = invocation.outputFormat; @@ -64,14 +51,12 @@ export const claudeHandler: AgentHandler = { const conversation = query({ prompt: invocation.userPrompt, - options: queryOptions as Parameters[0]["options"], + options: queryOptions, }); for await (const message of conversation) { if (message.type === "assistant") { - const text = extractTextFromContent( - message.message?.content ?? [], - ); + const text = extractTextFromContent(message.message?.content ?? []); if (text) { sender.sendAssistantMessage([text]); } @@ -90,8 +75,10 @@ export const claudeHandler: AgentHandler = { "Expected structured_output in result but got none", ); } - structuredOutput = - message.structured_output as Record; + structuredOutput = message.structured_output as Record< + string, + unknown + >; } else { if (message.structured_output != null) { throw new Error( diff --git a/packages/runtimeuse/src/command-handler.test.ts b/packages/runtimeuse/src/command-handler.test.ts index 295517f..cb4e27a 100644 --- a/packages/runtimeuse/src/command-handler.test.ts +++ b/packages/runtimeuse/src/command-handler.test.ts @@ -147,6 +147,38 @@ describe("CommandHandler", () => { ); }); + it("merges command env on top of process.env", () => { + const handler = createHandler({ + command: "echo", + cwd: "/work", + env: { MY_VAR: "hello", OTHER: "world" }, + }); + + handler.execute(); + + expect(exec).toHaveBeenCalledWith( + "echo", + expect.objectContaining({ + env: { ...process.env, MY_VAR: "hello", OTHER: "world" }, + }), + expect.any(Function), + ); + }); + + it("uses only process.env when command env is undefined", () => { + const handler = createHandler({ command: "echo", cwd: "/work" }); + + handler.execute(); + + expect(exec).toHaveBeenCalledWith( + "echo", + expect.objectContaining({ + env: { ...process.env }, + }), + expect.any(Function), + ); + }); + it("passes abort signal to exec", () => { const ac = new AbortController(); const handler = createHandler({ command: "sleep" }, ac); diff --git a/packages/runtimeuse/src/command-handler.ts b/packages/runtimeuse/src/command-handler.ts index 913ba2d..b6d2b09 100644 --- a/packages/runtimeuse/src/command-handler.ts +++ b/packages/runtimeuse/src/command-handler.ts @@ -48,7 +48,7 @@ class CommandHandler { this.command.command, { cwd: this.command.cwd ?? process.cwd(), - env: process.env, + env: { ...process.env, ...this.command.env }, signal: this.abortController.signal, }, (error, stdout, stderr) => { diff --git a/packages/runtimeuse/src/invocation-runner.test.ts b/packages/runtimeuse/src/invocation-runner.test.ts index 7e1889b..fcbd652 100644 --- a/packages/runtimeuse/src/invocation-runner.test.ts +++ b/packages/runtimeuse/src/invocation-runner.test.ts @@ -289,6 +289,66 @@ describe("InvocationRunner", () => { ); }); + it("passes agent_env to handler as env", async () => { + mockHandlerRun.mockResolvedValueOnce({ + type: "text", + text: "done", + } as AgentResult); + const { runner } = createRunner({ + agent_env: { API_KEY: "secret", MODE: "test" }, + output_format_json_schema_str: undefined, + }); + + await runner.run({ + ...BASE_INVOCATION_MESSAGE, + agent_env: { API_KEY: "secret", MODE: "test" }, + output_format_json_schema_str: undefined, + }); + + expect(mockHandlerRun).toHaveBeenCalledWith( + expect.objectContaining({ + env: { API_KEY: "secret", MODE: "test" }, + }), + expect.any(Object), + ); + }); + + it("passes undefined env when agent_env is not set", async () => { + mockHandlerRun.mockResolvedValueOnce({ + type: "text", + text: "done", + } as AgentResult); + const { runner } = createRunner(); + + await runner.run({ + ...BASE_INVOCATION_MESSAGE, + output_format_json_schema_str: undefined, + }); + + expect(mockHandlerRun).toHaveBeenCalledWith( + expect.objectContaining({ + env: undefined, + }), + expect.any(Object), + ); + }); + + it("passes pre-command env through to CommandHandler", async () => { + const { runner, message } = createRunner({ + pre_agent_invocation_commands: [ + { command: "setup", cwd: "/app", env: { SETUP_VAR: "1" } }, + ], + }); + + await runner.run(message); + + expect(CommandHandler).toHaveBeenCalledWith( + expect.objectContaining({ + command: { command: "setup", cwd: "/app", env: { SETUP_VAR: "1" } }, + }), + ); + }); + it("builds command handlers for each configured command", async () => { const { runner, message } = createRunner({ pre_agent_invocation_commands: [ @@ -413,6 +473,22 @@ describe("InvocationRunner.runCommandsOnly", () => { }); }); + it("passes command env through to CommandHandler", async () => { + const { runner, message } = createCommandRunner({ + commands: [ + { command: "echo hello", cwd: "/app", env: { FOO: "bar" } }, + ], + }); + + await runner.runCommandsOnly(message); + + expect(CommandHandler).toHaveBeenCalledWith( + expect.objectContaining({ + command: { command: "echo hello", cwd: "/app", env: { FOO: "bar" } }, + }), + ); + }); + it("downloads runtime environment before running commands", async () => { const events: string[] = []; mockDownload.mockImplementation(async () => { diff --git a/packages/runtimeuse/src/invocation-runner.ts b/packages/runtimeuse/src/invocation-runner.ts index 5b86c74..b4f4912 100644 --- a/packages/runtimeuse/src/invocation-runner.ts +++ b/packages/runtimeuse/src/invocation-runner.ts @@ -5,6 +5,7 @@ import type { CommandExecutionResultItem, OutgoingMessage, RuntimeEnvironmentDownloadable, + Command, } from "./types.js"; import type { Logger } from "./logger.js"; import CommandHandler from "./command-handler.js"; @@ -30,7 +31,11 @@ export class InvocationRunner { const { handler, logger, abortController, send } = this.config; await this.downloadRuntimeEnvironment(message.pre_agent_downloadables); - await this.runCommands(message.pre_agent_invocation_commands, "pre-agent", message.secrets_to_redact); + await this.runCommands( + message.pre_agent_invocation_commands, + "pre-agent", + message.secrets_to_redact, + ); const sender = this.createSender(); const outputFormat = message.output_format_json_schema_str @@ -46,6 +51,7 @@ export class InvocationRunner { userPrompt: message.user_prompt, outputFormat, model: message.model, + env: message.agent_env, secrets: message.secrets_to_redact, signal: abortController.signal, logger, @@ -82,19 +88,26 @@ export class InvocationRunner { const results: CommandExecutionResultItem[] = []; for (const command of message.commands) { - await this.runCommandAndCollect(command, message.secrets_to_redact, results); + await this.runCommandAndCollect( + command, + message.secrets_to_redact, + results, + ); } const resultMessage: OutgoingMessage = { message_type: "command_execution_result_message", results, }; - logger.log("Sending command execution result:", JSON.stringify(resultMessage)); + logger.log( + "Sending command execution result:", + JSON.stringify(resultMessage), + ); send(resultMessage); } private async runCommandAndCollect( - command: { command: string; cwd?: string }, + command: Command, secrets: string[], results: CommandExecutionResultItem[], ): Promise { diff --git a/packages/runtimeuse/src/types.ts b/packages/runtimeuse/src/types.ts index 433edac..8fe81b7 100644 --- a/packages/runtimeuse/src/types.ts +++ b/packages/runtimeuse/src/types.ts @@ -1,6 +1,7 @@ interface Command { command: string; cwd?: string; + env?: Record; } interface RuntimeEnvironmentDownloadable { @@ -13,6 +14,7 @@ interface InvocationMessage { source_id?: string; system_prompt: string; user_prompt: string; + agent_env?: Record; secrets_to_redact: string[]; output_format_json_schema_str?: string; model: string;