Skip to content

Commit ea87edc

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
fix: Remove sandbox_hostname from JWT payload.
PiperOrigin-RevId: 913779020
1 parent 4ba222b commit ea87edc

1 file changed

Lines changed: 2 additions & 17 deletions

File tree

vertexai/_genai/sandboxes.py

Lines changed: 2 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -870,19 +870,13 @@ def delete(
870870
def generate_access_token(
871871
self,
872872
service_account_email: str,
873-
sandbox_hostname: str,
874-
port: str = "8080",
875873
timeout: int = 3600,
876874
) -> str:
877875
"""Signs a JWT with a Google Cloud service account.
878876
879877
Args:
880878
service_account_email (str):
881879
Required. The email of the service account to use for signing.
882-
sandbox_hostname (str):
883-
Required. The hostname of the sandbox to generate a token for.
884-
port (str):
885-
Optional. The port to use for the token. Defaults to "8080".
886880
timeout (int):
887881
Optional. The timeout in seconds for the token. Defaults to 3600.
888882
@@ -891,15 +885,13 @@ def generate_access_token(
891885
"""
892886
client = iam_credentials_v1.IAMCredentialsClient()
893887
name = f"projects/-/serviceAccounts/{service_account_email}"
894-
custom_claims = {"hostname": sandbox_hostname, "port": port}
895888
payload = {
896889
"iat": int(time.time()),
897890
"exp": int(time.time()) + timeout,
898891
"iss": service_account_email,
899892
"sub": service_account_email,
900893
"nonce": secrets.randbelow(1000000000) + 1,
901894
"aud": "https://aiplatform.googleapis.com/", # default audience for sandbox proxy
902-
**custom_claims,
903895
}
904896
request = iam_credentials_v1.SignJwtRequest(
905897
name=name,
@@ -1010,9 +1002,7 @@ def generate_browser_ws_headers(
10101002
else:
10111003
raise ValueError("Load balancer hostname or ip is not available.")
10121004

1013-
http_access_token = self.generate_access_token(
1014-
service_account_email, connection_info.load_balancer_hostname, port, timeout
1015-
)
1005+
http_access_token = self.generate_access_token(service_account_email, timeout)
10161006
response = self.send_command(
10171007
http_method="GET",
10181008
access_token=http_access_token,
@@ -1027,12 +1017,7 @@ def generate_browser_ws_headers(
10271017
ws_url = ws_base_url + "/" + ws_path
10281018

10291019
# port 9222 is the default port for the browser websocket endpoint.
1030-
ws_access_token = self.generate_access_token(
1031-
service_account_email,
1032-
connection_info.load_balancer_hostname,
1033-
"9222",
1034-
timeout,
1035-
)
1020+
ws_access_token = self.generate_access_token(service_account_email, timeout)
10361021

10371022
routing_token = connection_info.routing_token
10381023

0 commit comments

Comments
 (0)