diff --git a/src/hai_agents_local/workstation.py b/src/hai_agents_local/workstation.py index 1d03225..6b78a1d 100644 --- a/src/hai_agents_local/workstation.py +++ b/src/hai_agents_local/workstation.py @@ -9,6 +9,7 @@ import subprocess import sys import sysconfig +from collections.abc import Mapping from pathlib import Path from typing import TYPE_CHECKING @@ -36,9 +37,11 @@ def __init__( api_key: TokenSource, base_url: str | None = None, session_id: str | None = None, + environment_variables: Mapping[str, str] | None = None, ) -> None: super().__init__(environment_id, api_key=api_key, base_url=base_url, session_id=session_id) self.workspace = Path(workspace).expanduser() if workspace else Path.home() / "hai" / self.session_id + self.environment_variables = dict(environment_variables or {}) def preflight(self) -> None: if sys.platform == "darwin": @@ -58,6 +61,7 @@ def create_driver(self) -> ManagedCodeSandboxInterface: return LocalCodeSandbox( str(self.workspace), environment_variables={ + **self.environment_variables, "COORDINATE_SYSTEM": "0-1000", "CUA_CURSOR": os.environ.get("CUA_CURSOR", "off"), "CUA_SESSION": self.session_id, diff --git a/tests/test_local.py b/tests/test_local.py index 8bc829d..06a293b 100644 --- a/tests/test_local.py +++ b/tests/test_local.py @@ -731,17 +731,23 @@ def test_workstation_serves_a_shell_in_its_session_folder_with_the_cli_commands( pytest.importorskip("hai_drivers.code_sandbox.local.driver") monkeypatch.setenv("HOME", str(tmp_path)) monkeypatch.delenv("CUA_CURSOR", raising=False) - bridge = WorkstationBridge(api_key="k") + bridge = WorkstationBridge( + api_key="k", + environment_variables={"RUN_CONFIG": "/run/config.json", "CUA_SESSION": "other", "PATH": "/nowhere"}, + ) bridge._driver = bridge.create_driver() try: result, error = bridge._dispatch( "execute", - {"command": "pwd -P; echo $COORDINATE_SYSTEM $CUA_SESSION $CUA_CURSOR; command -v desk web cua click"}, + { + "command": "pwd -P; echo $COORDINATE_SYSTEM $CUA_SESSION $CUA_CURSOR $RUN_CONFIG;" + " command -v desk web cua click" + }, ) assert error is None and result["exit_code"] == 0 out = result["stdout"].split() assert out[0] == str((tmp_path / "hai" / bridge.session_id).resolve()) - assert out[1:4] == ["0-1000", bridge.session_id, "off"] and len(out) == 8 + assert out[1:5] == ["0-1000", bridge.session_id, "off", "/run/config.json"] and len(out) == 9 finally: bridge._driver.close()