feat: Add psc_endpoint parameter to Sandboxes.send_command and Sandboxes.generate_browser_ws_headers so VPC-SC sandboxes with ingressControlConfig.enablePrivateServiceConnect=true (which populate only connection_info.service_attachment, never load_balancer_hostname / load_balancer_ip) can route data-plane traffic through a customer-provisioned PSC endpoint. Existing non-VPC-SC callers are unaffected; the parameter defaults to None and the load-balancer branches remain the primary path.

PiperOrigin-RevId: 981254234

Author: vertex-sdk-bot

agentplatform/_genai/sandboxes.py

@@ -1362,6 +1362,7 @@ def send_command(
         query_params: Optional[dict[str, object]] = None,
         headers: Optional[dict[str, str]] = None,
         request_dict: Optional[dict[str, object]] = None,
+        psc_endpoint: Optional[str] = None,
     ) -> genai_types.HttpResponse:
         """Sends a command to the sandbox.
 
@@ -1382,6 +1383,16 @@ def send_command(
                 Optional. The headers to include in the command.
             request_dict (dict[str, object]):
                 Optional. The request body to include in the command.
+            psc_endpoint (str):
+                Optional. Host (RFC1918 IP or resolvable internal hostname) of the
+                customer-side PSC endpoint forwarding rule that targets
+                ``connection_info.service_attachment``. Needed for VPC-SC sandboxes
+                (``ingressControlConfig.enablePrivateServiceConnect=true``) where
+                ``connection_info.load_balancer_hostname`` and ``load_balancer_ip``
+                are both unset because ingress goes through PSC in the customer's
+                VPC rather than a Google-managed load balancer. When provided, it
+                is used as the data-plane host in place of the load balancer
+                address.
 
         Returns:
             genai_types.HttpResponse: The response from the sandbox.
@@ -1391,12 +1402,21 @@ def send_command(
         connection_info = sandbox_environment.connection_info
         if not connection_info:
             raise ValueError("Connection info is not available.")
-        if connection_info.load_balancer_hostname:
+        if psc_endpoint:
+            # VPC-SC path: caller has provisioned a PSC endpoint against
+            # connection_info.service_attachment and passes its host here.
+            endpoint = "https://" + psc_endpoint
+        elif connection_info.load_balancer_hostname:
             endpoint = "https://" + connection_info.load_balancer_hostname
         elif connection_info.load_balancer_ip:
             endpoint = "http://" + connection_info.load_balancer_ip
         else:
-            raise ValueError("Load balancer hostname or ip is not available.")
+            raise ValueError(
+                "No data-plane endpoint available. Non-VPC-SC sandboxes populate"
+                " connection_info.load_balancer_hostname; VPC-SC sandboxes require"
+                " the caller to pass psc_endpoint (the host of the forwarding"
+                " rule targeting connection_info.service_attachment)."
+            )
 
         routing_token = connection_info.routing_token
         if not routing_token:
@@ -1437,6 +1457,7 @@ def generate_browser_ws_headers(
         service_account_email: str,
         port: str = "8080",
         timeout: int = 3600,
+        psc_endpoint: Optional[str] = None,
     ) -> tuple[str, dict[str, str]]:
         """Generates the websocket upgrade headers for the browser.
 
@@ -1450,6 +1471,15 @@ def generate_browser_ws_headers(
                 Defaults to "8080". This should be one of the ports specified during template creation.
             timeout (int):
                 Optional. The timeout in seconds for the token. Defaults to 3600.
+            psc_endpoint (str):
+                Optional. Host (RFC1918 IP or resolvable internal hostname) of the
+                customer-side PSC endpoint forwarding rule that targets
+                ``connection_info.service_attachment``. Needed for VPC-SC
+                sandboxes; see ``send_command`` for details. When provided, it is
+                used as the websocket host in place of the load balancer address,
+                and is propagated to the internal ``send_command`` call that
+                fetches the CDP endpoint.
+
         Returns:
             tuple[str, dict[str, str]]: A tuple containing the websocket URL and
             the headers for websocket upgrade.
@@ -1458,12 +1488,19 @@ def generate_browser_ws_headers(
             raise ValueError("Connection info is not available.")
 
         connection_info = sandbox_environment.connection_info
-        if connection_info.load_balancer_hostname:
+        if psc_endpoint:
+            ws_base_url = "wss://" + psc_endpoint
+        elif connection_info.load_balancer_hostname:
             ws_base_url = "wss://" + connection_info.load_balancer_hostname
         elif connection_info.load_balancer_ip:
             ws_base_url = "ws://" + connection_info.load_balancer_ip
         else:
-            raise ValueError("Load balancer hostname or ip is not available.")
+            raise ValueError(
+                "No data-plane endpoint available. Non-VPC-SC sandboxes populate"
+                " connection_info.load_balancer_hostname; VPC-SC sandboxes require"
+                " the caller to pass psc_endpoint (the host of the forwarding"
+                " rule targeting connection_info.service_attachment)."
+            )
 
         http_access_token = self.generate_access_token(service_account_email, timeout)
         response = self.send_command(
@@ -1472,6 +1509,7 @@ def generate_browser_ws_headers(
             sandbox_environment=sandbox_environment,
             port=port,
             path="/cdp_ws_endpoint",
+            psc_endpoint=psc_endpoint,
         )
         if not response:
             raise ValueError("Failed to get the websocket endpoint.")

tests/unit/agentplatform/genai/test_sandbox.py

@@ -164,6 +164,106 @@ def test_generate_browser_ws_headers(
             == "v1.stream, test_token, test_routing_token, 9222"
         )
 
+    @mock.patch.object(sandboxes.requests, "request")
+    def test_send_command_vpcsc_uses_psc_endpoint(self, mock_request):
+        mock_sandbox = mock.Mock()
+        mock_sandbox.connection_info.load_balancer_hostname = None
+        mock_sandbox.connection_info.load_balancer_ip = None
+        mock_sandbox.connection_info.routing_token = "test_routing_token"
+        mock_response = mock.Mock()
+        mock_response.text = "{}"
+        mock_response.headers = {}
+        mock_request.return_value = mock_response
+
+        self.client.sandboxes.send_command(
+            http_method="GET",
+            access_token="test_token",
+            sandbox_environment=mock_sandbox,
+            port="9000",
+            path="test/path",
+            psc_endpoint="10.0.0.10",
+        )
+
+        args, kwargs = mock_request.call_args
+        assert args[0] == "GET"
+        assert args[1] == "https://10.0.0.10/test/path"
+        assert kwargs["headers"]["Authorization"] == "Bearer test_token"
+        assert kwargs["headers"]["X-Sandbox-Routing-Token"] == "test_routing_token"
+        # X-Sandbox-Port must flow through even on the VPC-SC path; the
+        # reverse proxy behind the customer's PSC endpoint uses it to route
+        # to the correct sandbox container port.
+        assert kwargs["headers"]["X-Sandbox-Port"] == "9000"
+
+    @mock.patch.object(sandboxes.requests, "request")
+    def test_send_command_vpcsc_accepts_hostname_as_psc_endpoint(self, mock_request):
+        """psc_endpoint accepts a resolvable hostname, not only an IP."""
+        mock_sandbox = mock.Mock()
+        mock_sandbox.connection_info.load_balancer_hostname = None
+        mock_sandbox.connection_info.load_balancer_ip = None
+        mock_sandbox.connection_info.routing_token = "test_routing_token"
+        mock_response = mock.Mock()
+        mock_response.text = "{}"
+        mock_response.headers = {}
+        mock_request.return_value = mock_response
+
+        self.client.sandboxes.send_command(
+            http_method="GET",
+            access_token="test_token",
+            sandbox_environment=mock_sandbox,
+            path="test/path",
+            psc_endpoint="psc-endpoint.internal.example.com",
+        )
+
+        args, _ = mock_request.call_args
+        assert args[1] == "https://psc-endpoint.internal.example.com/test/path"
+
+    def test_send_command_raises_when_no_endpoint_available(self):
+        mock_sandbox = mock.Mock()
+        mock_sandbox.connection_info.load_balancer_hostname = None
+        mock_sandbox.connection_info.load_balancer_ip = None
+        mock_sandbox.connection_info.routing_token = "test_routing_token"
+
+        with pytest.raises(ValueError, match="psc_endpoint"):
+            self.client.sandboxes.send_command(
+                http_method="GET",
+                access_token="test_token",
+                sandbox_environment=mock_sandbox,
+                path="test/path",
+            )
+
+    @mock.patch.object(sandboxes.Sandboxes, "generate_access_token")
+    @mock.patch.object(sandboxes.requests, "request")
+    def test_generate_browser_ws_headers_vpcsc_uses_psc_endpoint(
+        self, mock_request, mock_generate_access_token
+    ):
+        mock_generate_access_token.return_value = "test_token"
+
+        mock_sandbox = mock.Mock()
+        mock_sandbox.connection_info.load_balancer_hostname = None
+        mock_sandbox.connection_info.load_balancer_ip = None
+        mock_sandbox.connection_info.routing_token = "test_routing_token"
+        mock_response = mock.Mock()
+        mock_response.text = '{"endpoint": "test/endpoint"}'
+        mock_response.headers = {}
+        mock_request.return_value = mock_response
+
+        ws_url, headers = self.client.sandboxes.generate_browser_ws_headers(
+            sandbox_environment=mock_sandbox,
+            service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL,
+            timeout=3600,
+            psc_endpoint="10.0.0.10",
+        )
+        assert ws_url == "wss://10.0.0.10/test/endpoint"
+        assert (
+            headers["Sec-WebSocket-Protocol"]
+            == "v1.stream, test_token, test_routing_token, 9222"
+        )
+        # The delegated send_command call that fetches the CDP endpoint must
+        # also target the PSC endpoint host, not the (empty) load balancer.
+        req_args, req_kwargs = mock_request.call_args
+        assert req_args[1] == "https://10.0.0.10/cdp_ws_endpoint"
+        assert req_kwargs["headers"]["X-Sandbox-Routing-Token"] == "test_routing_token"
+
     @mock.patch.object(sandboxes.Sandboxes, "_create")
     def test_create_with_shell_environment_and_existing_template(self, mock_create):
         mock_operation = mock.Mock()

vertexai/_genai/sandboxes.py

@@ -937,6 +937,7 @@ def send_command(
         query_params: Optional[dict[str, object]] = None,
         headers: Optional[dict[str, str]] = None,
         request_dict: Optional[dict[str, object]] = None,
+        psc_endpoint: Optional[str] = None,
     ) -> genai_types.HttpResponse:
         """Sends a command to the sandbox.
 
@@ -957,6 +958,16 @@ def send_command(
                 Optional. The headers to include in the command.
             request_dict (dict[str, object]):
                 Optional. The request body to include in the command.
+            psc_endpoint (str):
+                Optional. Host (RFC1918 IP or resolvable internal hostname) of the
+                customer-side PSC endpoint forwarding rule that targets
+                ``connection_info.service_attachment``. Needed for VPC-SC sandboxes
+                (``ingressControlConfig.enablePrivateServiceConnect=true``) where
+                ``connection_info.load_balancer_hostname`` and ``load_balancer_ip``
+                are both unset because ingress goes through PSC in the customer's
+                VPC rather than a Google-managed load balancer. When provided, it
+                is used as the data-plane host in place of the load balancer
+                address.
 
         Returns:
             genai_types.HttpResponse: The response from the sandbox.
@@ -966,12 +977,21 @@ def send_command(
         connection_info = sandbox_environment.connection_info
         if not connection_info:
             raise ValueError("Connection info is not available.")
-        if connection_info.load_balancer_hostname:
+        if psc_endpoint:
+            # VPC-SC path: caller has provisioned a PSC endpoint against
+            # connection_info.service_attachment and passes its host here.
+            endpoint = "https://" + psc_endpoint
+        elif connection_info.load_balancer_hostname:
             endpoint = "https://" + connection_info.load_balancer_hostname
         elif connection_info.load_balancer_ip:
             endpoint = "http://" + connection_info.load_balancer_ip
         else:
-            raise ValueError("Load balancer hostname or ip is not available.")
+            raise ValueError(
+                "No data-plane endpoint available. Non-VPC-SC sandboxes populate"
+                " connection_info.load_balancer_hostname; VPC-SC sandboxes require"
+                " the caller to pass psc_endpoint (the host of the forwarding"
+                " rule targeting connection_info.service_attachment)."
+            )
 
         routing_token = connection_info.routing_token
         if not routing_token:
@@ -1012,6 +1032,7 @@ def generate_browser_ws_headers(
         service_account_email: str,
         port: str = "8080",
         timeout: int = 3600,
+        psc_endpoint: Optional[str] = None,
     ) -> tuple[str, dict[str, str]]:
         """Generates the websocket upgrade headers for the browser.
 
@@ -1025,6 +1046,15 @@ def generate_browser_ws_headers(
                 Defaults to "8080". This should be one of the ports specified during template creation.
             timeout (int):
                 Optional. The timeout in seconds for the token. Defaults to 3600.
+            psc_endpoint (str):
+                Optional. Host (RFC1918 IP or resolvable internal hostname) of the
+                customer-side PSC endpoint forwarding rule that targets
+                ``connection_info.service_attachment``. Needed for VPC-SC
+                sandboxes; see ``send_command`` for details. When provided, it is
+                used as the websocket host in place of the load balancer address,
+                and is propagated to the internal ``send_command`` call that
+                fetches the CDP endpoint.
+
         Returns:
             tuple[str, dict[str, str]]: A tuple containing the websocket URL and
             the headers for websocket upgrade.
@@ -1033,12 +1063,19 @@ def generate_browser_ws_headers(
             raise ValueError("Connection info is not available.")
 
         connection_info = sandbox_environment.connection_info
-        if connection_info.load_balancer_hostname:
+        if psc_endpoint:
+            ws_base_url = "wss://" + psc_endpoint
+        elif connection_info.load_balancer_hostname:
             ws_base_url = "wss://" + connection_info.load_balancer_hostname
         elif connection_info.load_balancer_ip:
             ws_base_url = "ws://" + connection_info.load_balancer_ip
         else:
-            raise ValueError("Load balancer hostname or ip is not available.")
+            raise ValueError(
+                "No data-plane endpoint available. Non-VPC-SC sandboxes populate"
+                " connection_info.load_balancer_hostname; VPC-SC sandboxes require"
+                " the caller to pass psc_endpoint (the host of the forwarding"
+                " rule targeting connection_info.service_attachment)."
+            )
 
         http_access_token = self.generate_access_token(service_account_email, timeout)
         response = self.send_command(
@@ -1047,6 +1084,7 @@ def generate_browser_ws_headers(
             sandbox_environment=sandbox_environment,
             port=port,
             path="/cdp_ws_endpoint",
+            psc_endpoint=psc_endpoint,
         )
         if not response:
             raise ValueError("Failed to get the websocket endpoint.")