No public description

PiperOrigin-RevId: 945254181
Change-Id: Ib2951de932f2d61445df568aa7961089be647f34

Author: Antigravity Team

google/antigravity/connections/local/litert_connection_config.py

@@ -160,4 +160,5 @@ def create_strategy(
         skills_paths=self.skills_paths,
         mcp_servers=self.mcp_servers,
         subagents=self.subagents,
+        env=self.env,
     )

google/antigravity/connections/local/local_connection.py

@@ -173,11 +173,13 @@ def __init__(
       tool_runner: t_runner.ToolRunner | None = None,
       hook_runner: h_runner.HookRunner | None = None,
       initial_history: Sequence[types.Step] | None = None,
+      env: dict[str, str] | None = None,
   ):
     self._hook_runner = hook_runner
     self._process = process
     self._ws = ws
     self._tool_runner = tool_runner
+    self._env = env
     self.__initial_history = initial_history or []
     self._client_cancelled = False
     self._is_receiving = False
@@ -623,6 +625,7 @@ def __init__(
       workspaces: list[str] | None = None,
       app_data_dir: str | None = None,
       mcp_servers: Sequence[types.McpServerConfig] | None = None,
+      env: dict[str, str] | None = None,
       subagents: list[types.SubagentConfig] | None = None,
   ):
     """Initializes the instance.
@@ -639,6 +642,7 @@ def __init__(
       workspaces: Optional list of workspace paths.
       app_data_dir: Optional directory for harness app data.
       mcp_servers: Optional sequence of MCP server configurations.
+      env: Optional dictionary of custom environment variables.
       subagents: Optional list of static subagent configurations.
     """
     self._binary_path = _get_default_binary_path()
@@ -648,6 +652,7 @@ def __init__(
     self._mcp_servers = mcp_servers or []
     self._models: list[types.ModelTarget] = models or []
     self._skills_paths = skills_paths
+    self._env = env
 
     # Normalize str shorthand to SystemInstructions model.
     self._system_instructions: types.SystemInstructions | None = None
@@ -974,16 +979,23 @@ async def __aenter__(self) -> None:
         version=sdk_version,
         language_version=platform.python_version(),
     )
+    env_map = (
+        {str(k): str(v) for k, v in self._env.items()} if self._env else {}
+    )
     input_config = localharness_pb2.InputConfig(
         storage_directory=self._save_dir or "",
         client_info=client_info_proto,
+        env=env_map,
     )
 
+    merged_env = {**os.environ, **env_map} if self._env is not None else None
+
     process = subprocess.Popen(
         [self._binary_path],
         stdin=subprocess.PIPE,
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
+        env=merged_env,
     )
 
     serialized = input_config.SerializeToString()
@@ -1058,6 +1070,7 @@ async def __aenter__(self) -> None:
         tool_runner=self._tool_runner,
         hook_runner=self._hook_runner,
         initial_history=initial_history,
+        env=self._env,
     )
     self._connection._start_stderr_reader(process.stderr)
 

google/antigravity/connections/local/local_connection_config.py

@@ -97,6 +97,13 @@ class BaseLocalAgentConfig(connection.AgentConfig):
   policies: list[policy.Policy] = pydantic.Field(
       default_factory=policy.confirm_run_command
   )
+  env: dict[str, str] | None = pydantic.Field(
+      default=None,
+      description=(
+          "Optional custom environment variables for the localharness"
+          " subprocess."
+      ),
+  )
   workspaces: list[str] = pydantic.Field(default_factory=lambda: [os.getcwd()])
 
   @pydantic.field_validator("app_data_dir")
@@ -171,6 +178,7 @@ def __init__(  # pylint: disable=super-init-not-called
       mcp_servers: list[types.McpServerConfig] | None = None,
       subagents: list[types.SubagentConfig] | None = None,
       workspaces: list[str] | None = None,
+      env: dict[str, str] | None = None,
       conversation_id: str | None = None,
       save_dir: str | None = None,
       app_data_dir: str | None = None,
@@ -299,5 +307,6 @@ def create_strategy(
         app_data_dir=self.app_data_dir,
         skills_paths=self.skills_paths,
         mcp_servers=self.mcp_servers,
+        env=self.env,
         subagents=self.subagents,
     )

google/antigravity/connections/local/local_connection_test.py

@@ -1840,6 +1840,89 @@ async def test_accepts_env_var_api_key(self, mock_popen):
       async with strategy:
         pass
 
+  @mock.patch.dict(
+      "os.environ",
+      {"GEMINI_API_KEY": "env-key", "SYS_VAR": "sys_val"},
+      clear=True,
+  )
+  @mock.patch("subprocess.Popen")
+  async def test_passes_custom_env_to_popen_and_input_config(self, mock_popen):
+    """Verifies custom env dict is merged into Popen env and InputConfig."""
+    mock_proc = mock.MagicMock()
+    mock_proc.stdin = mock.MagicMock()
+    mock_proc.stdout = mock.MagicMock()
+    mock_proc.stderr = mock.MagicMock()
+    mock_proc.stdout.read.return_value = b""
+    mock_popen.return_value = mock_proc
+
+    custom_env = {"MY_CUSTOM_VAR": "hello_env"}
+    strategy = self._make_strategy(env=custom_env)
+
+    with self.assertRaises(RuntimeError):
+      async with strategy:
+        pass
+
+    mock_popen.assert_called_once()
+    _, kwargs = mock_popen.call_args
+    expected_env = {
+        "GEMINI_API_KEY": "env-key",
+        "SYS_VAR": "sys_val",
+        "MY_CUSTOM_VAR": "hello_env",
+    }
+    self.assertEqual(kwargs.get("env"), expected_env)
+
+    mock_proc.stdin.write.assert_called_once()
+    written_bytes = mock_proc.stdin.write.call_args[0][0]
+    parsed_config = localharness_pb2.InputConfig()
+    parsed_config.ParseFromString(written_bytes[4:])
+    self.assertEqual(dict(parsed_config.env), custom_env)
+
+  @mock.patch.dict(
+      "os.environ",
+      {"GEMINI_API_KEY": "env-key"},
+      clear=True,
+  )
+  @mock.patch("subprocess.Popen")
+  async def test_passes_non_string_env_coerced_to_strings(self, mock_popen):
+    """Verifies non-string env keys/values are coerced to strings for Popen and InputConfig."""
+    mock_proc = mock.MagicMock()
+    mock_proc.stdin = mock.MagicMock()
+    mock_proc.stdout = mock.MagicMock()
+    mock_proc.stderr = mock.MagicMock()
+    mock_proc.stdout.read.return_value = b""
+    mock_popen.return_value = mock_proc
+
+    custom_env = {"INT_KEY": 123, 1: "val", "BOOL_KEY": True}
+    strategy = self._make_strategy(env=custom_env)
+
+    with self.assertRaises(RuntimeError):
+      async with strategy:
+        pass
+
+    mock_popen.assert_called_once()
+    _, kwargs = mock_popen.call_args
+    expected_env = {
+        "GEMINI_API_KEY": "env-key",
+        "INT_KEY": "123",
+        "1": "val",
+        "BOOL_KEY": "True",
+    }
+    self.assertEqual(kwargs.get("env"), expected_env)
+
+    mock_proc.stdin.write.assert_called_once()
+    written_bytes = mock_proc.stdin.write.call_args[0][0]
+    parsed_config = localharness_pb2.InputConfig()
+    parsed_config.ParseFromString(written_bytes[4:])
+    self.assertEqual(
+        dict(parsed_config.env),
+        {"INT_KEY": "123", "1": "val", "BOOL_KEY": "True"},
+    )
+
+  def test_config_env_defaults_to_none(self):
+    """Verifies that LocalAgentConfig.env is None by default."""
+    config = local_connection_config.LocalAgentConfig()
+    self.assertIsNone(config.env)
+
   @mock.patch.dict("os.environ", {}, clear=True)
   @mock.patch("subprocess.Popen")
   async def test_accepts_models_api_key(self, mock_popen):
@@ -3198,6 +3281,22 @@ def test_create_strategy(self):
     self.assertLen(text_models, 1)
     self.assertEqual(text_models[0].name, "gemini-2.5-pro")
 
+  def test_create_strategy_passes_env(self):
+    config = local_connection_config.LocalAgentConfig(
+        env={"CUSTOM_KEY": "CUSTOM_VAL"},
+    )
+    mock_tool_runner = mock.create_autospec(
+        tool_runner.ToolRunner, instance=True
+    )
+    mock_hook_runner = mock.create_autospec(
+        hook_runner.HookRunner, instance=True
+    )
+    strategy = config.create_strategy(
+        tool_runner=mock_tool_runner,
+        hook_runner=mock_hook_runner,
+    )
+    self.assertEqual(strategy._env, {"CUSTOM_KEY": "CUSTOM_VAL"})
+
   def test_merge_models_only_defaults(self):
     config = local_connection_config.LocalAgentConfig()
     self.assertLen(config.models, 2)

google/antigravity/connections/local/local_openai_connection_config.py

@@ -113,4 +113,5 @@ def create_strategy(
         skills_paths=self.skills_paths,
         mcp_servers=self.mcp_servers,
         subagents=self.subagents,
+        env=self.env,
     )