Skip to content
Open
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
15 changes: 13 additions & 2 deletions agentplatform/_genai/sandboxes.py
Original file line number Diff line number Diff line change
Expand Up @@ -1405,7 +1405,14 @@ def send_command(
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
# The forwarding rule targets the router's Service Attachment
# directly, and the router pod serves plain HTTP on `port` (the
# router NetworkPolicy allows ingress only on that TCP port).
# Unlike the non-VPC-SC load_balancer_hostname case, no
# Google-managed L7 LB terminates TLS or translates ports in
# front, so the URL must be http://host:port, not https://host
# on the default 443.
endpoint = f"http://{psc_endpoint}:{port}"
elif connection_info.load_balancer_hostname:
endpoint = "https://" + connection_info.load_balancer_hostname
elif connection_info.load_balancer_ip:
Expand Down Expand Up @@ -1489,7 +1496,11 @@ def generate_browser_ws_headers(

connection_info = sandbox_environment.connection_info
if psc_endpoint:
ws_base_url = "wss://" + psc_endpoint
# See send_command's psc_endpoint branch for the rationale: the
# PSC forwarding rule targets the router's Service Attachment
# directly (no L7 LB in front), so this must be plain ws:// with
# the explicit port, not wss:// on the default 443.
ws_base_url = f"ws://{psc_endpoint}:{port}"
elif connection_info.load_balancer_hostname:
ws_base_url = "wss://" + connection_info.load_balancer_hostname
elif connection_info.load_balancer_ip:
Expand Down
18 changes: 14 additions & 4 deletions tests/unit/agentplatform/genai/test_sandbox.py
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,11 @@ def test_send_command_vpcsc_uses_psc_endpoint(self, mock_request):

args, kwargs = mock_request.call_args
assert args[0] == "GET"
assert args[1] == "https://10.0.0.10/test/path"
# PSC-E ingress goes to the router pod directly (no Google-managed L7
# LB in front), and the router serves plain HTTP on the sandbox port.
# The URL must therefore be http://host:port, not https://host on the
# default 443.
assert args[1] == "http://10.0.0.10:9000/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
Expand Down Expand Up @@ -215,7 +219,11 @@ def test_send_command_vpcsc_accepts_hostname_as_psc_endpoint(self, mock_request)
)

args, _ = mock_request.call_args
assert args[1] == "https://psc-endpoint.internal.example.com/test/path"
# Default port "8080" is baked into the URL; scheme is plain http.
assert (
args[1]
== "http://psc-endpoint.internal.example.com:8080/test/path"
)

def test_send_command_raises_when_no_endpoint_available(self):
mock_sandbox = mock.Mock()
Expand Down Expand Up @@ -253,15 +261,17 @@ def test_generate_browser_ws_headers_vpcsc_uses_psc_endpoint(
timeout=3600,
psc_endpoint="10.0.0.10",
)
assert ws_url == "wss://10.0.0.10/test/endpoint"
# Plain ws:// with the explicit port for the same reason send_command
# uses http://host:port on this branch — see that method's docstring.
assert ws_url == "ws://10.0.0.10:8080/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_args[1] == "http://10.0.0.10:8080/cdp_ws_endpoint"
assert req_kwargs["headers"]["X-Sandbox-Routing-Token"] == "test_routing_token"

@mock.patch.object(sandboxes.Sandboxes, "_create")
Expand Down
Loading