fix: preserve OpenSandbox connection compatibility

Signed-off-by: Hemil Desai <hemild@nvidia.com>

Author: hemildesai

nemo_gym/sandbox/providers/opensandbox/provider.py

@@ -642,6 +642,8 @@ def _connection_config(
             kwargs["request_timeout"] = timedelta(seconds=request_timeout_s)
         if self._connection.use_server_proxy:
             kwargs["use_server_proxy"] = True
+            if self._connection.api_key is not None:
+                kwargs["headers"] = {"OPEN-SANDBOX-API-KEY": self._connection.api_key}
         if self._connection.keepalive_expiry_s is not None or self._connection.disable_connection_pooling:
             kwargs["transport"] = self._get_transport()
         config = ConnectionConfig(**kwargs)
@@ -1114,12 +1116,12 @@ async def endpoint(self, handle: SandboxHandle, port: int) -> SandboxEndpoint:
             # defaults match the connection that produced this endpoint.
             scheme = connection.get_base_url().split("://", 1)[0]
             endpoint_url = f"{scheme}://{endpoint_url.lstrip('/')}"
-        headers = dict(getattr(resolved, "headers", None) or {})
-        api_key = connection.headers.get("OPEN-SANDBOX-API-KEY")
-        if connection.use_server_proxy and api_key:
-            # Direct endpoints terminate at untrusted sandbox code; only proxy
-            # endpoints may receive server credentials.
-            headers.setdefault("OPEN-SANDBOX-API-KEY", api_key)
+        headers = dict(getattr(connection, "headers", None) or {})
+        headers.update(getattr(resolved, "headers", None) or {})
+        if not getattr(connection, "use_server_proxy", self._connection.use_server_proxy):
+            # Direct endpoints terminate at untrusted sandbox code and must
+            # never receive the management credential.
+            headers.pop("OPEN-SANDBOX-API-KEY", None)
         return SandboxEndpoint(endpoint=endpoint_url, headers=headers)
 
     def _command_retry_count(self) -> int:

tests/unit_tests/test_opensandbox_endpoint.py

@@ -116,7 +116,7 @@ def test_schemeless_endpoint_uses_the_sdk_resolved_url_and_key() -> None:
 def test_direct_mode_endpoint_does_not_inject_the_key() -> None:
     raw = _RawWithEndpoint(
         endpoint="http://pod.example:6000",
-        headers={},
+        headers={"OPEN-SANDBOX-API-KEY": "secret"},  # pragma: allowlist secret
         connection=_Connection(api_key="secret"),  # pragma: allowlist secret
     )
     resolved = asyncio.run(_provider().endpoint(_handle(raw), 6000))

tests/unit_tests/test_opensandbox_provider.py

@@ -365,6 +365,7 @@ def test_connection_config_and_image_policy(
         "protocol": "https",
         "request_timeout": timedelta(seconds=10),
         "use_server_proxy": True,
+        "headers": {"OPEN-SANDBOX-API-KEY": "key"},  # pragma: allowlist secret
     }
     assert config.headers == {"OPEN-SANDBOX-API-KEY": "key"}  # pragma: allowlist secret
     short_timeout_config = provider._connection_config(request_timeout_s=3)