From 39971ae242906c4f03833f53f7c7b2aadb772df4 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 15 Sep 2026 14:08:35 -0700 Subject: [PATCH] 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: 982031703 --- agentplatform/_genai/sandboxes.py | 46 +++++++- .../unit/agentplatform/genai/test_sandbox.py | 100 ++++++++++++++++++ vertexai/_genai/sandboxes.py | 46 +++++++- 3 files changed, 184 insertions(+), 8 deletions(-) diff --git a/agentplatform/_genai/sandboxes.py b/agentplatform/_genai/sandboxes.py index bce278816b..649cf3a743 100644 --- a/agentplatform/_genai/sandboxes.py +++ b/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.") diff --git a/tests/unit/agentplatform/genai/test_sandbox.py b/tests/unit/agentplatform/genai/test_sandbox.py index c38ba9d549..0e6c933a16 100644 --- a/tests/unit/agentplatform/genai/test_sandbox.py +++ b/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() diff --git a/vertexai/_genai/sandboxes.py b/vertexai/_genai/sandboxes.py index df582e74df..823c1d7995 100644 --- a/vertexai/_genai/sandboxes.py +++ b/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.")