Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 42 additions & 4 deletions agentplatform/_genai/sandboxes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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.
Expand All @@ -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:
Expand Down Expand Up @@ -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.

Expand All @@ -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.
Expand All @@ -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(
Expand All @@ -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.")
Expand Down
100 changes: 100 additions & 0 deletions tests/unit/agentplatform/genai/test_sandbox.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
46 changes: 42 additions & 4 deletions vertexai/_genai/sandboxes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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.
Expand All @@ -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:
Expand Down Expand Up @@ -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.

Expand All @@ -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.
Expand All @@ -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(
Expand All @@ -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.")
Expand Down
Loading