From 8933d1f2b9b15d60a156ff424d84602b5c86441b Mon Sep 17 00:00:00 2001 From: OhYee Date: Wed, 1 Apr 2026 15:28:30 +0800 Subject: [PATCH 1/4] feat(browser-api): enhance WebSocket URL generation with header support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add overloads to get_cdp_url and get_vnc_url methods to support returning authentication headers along with URLs. Introduce Config parameter for optional configuration overrides and improve URL assembly logic with proper tenant ID and recording parameters. The changes provide backward compatibility while enabling header-based authentication for WebSocket connections. feat(browser-api): 增强WebSocket URL生成以支持头部信息 为get_cdp_url和get_vnc_url方法添加重载以支持返回认证头部与URL。 引入Config参数用于可选配置覆盖，并改进URL组装逻辑，包含适当的租户ID和录制参数。 这些更改提供向后兼容性，同时启用基于头部的WebSocket连接认证。 Change-Id: Id927aa50801374160670c13b8a4dfab37d302ad9 Signed-off-by: OhYee --- .../sandbox/__aio_sandbox_async_template.py | 81 +++++++- .../__browser_sandbox_async_template.py | 79 +++++++- agentrun/sandbox/aio_sandbox.py | 81 +++++++- .../sandbox/api/__aio_data_async_template.py | 176 ++++++++++++----- .../api/__browser_data_async_template.py | 180 +++++++++++++----- agentrun/sandbox/api/aio_data.py | 176 ++++++++++++----- agentrun/sandbox/api/browser_data.py | 180 +++++++++++++----- agentrun/sandbox/browser_sandbox.py | 79 +++++++- agentrun/sandbox/custom_sandbox.py | 17 +- tests/unittests/sandbox/api/test_aio_data.py | 46 ++++- .../sandbox/api/test_browser_data.py | 47 ++++- .../unittests/sandbox/test_browser_sandbox.py | 4 +- .../unittests/sandbox/test_custom_sandbox.py | 38 ++-- 13 files changed, 918 insertions(+), 266 deletions(-) diff --git a/agentrun/sandbox/__aio_sandbox_async_template.py b/agentrun/sandbox/__aio_sandbox_async_template.py index eb7afc5..202529c 100644 --- a/agentrun/sandbox/__aio_sandbox_async_template.py +++ b/agentrun/sandbox/__aio_sandbox_async_template.py @@ -7,10 +7,11 @@ import asyncio import time # noqa: F401 -from typing import Optional +from typing import Dict, Literal, Optional, overload, Tuple, Union from agentrun.sandbox.api.aio_data import AioDataAPI from agentrun.sandbox.model import CodeLanguage, TemplateType +from agentrun.utils.config import Config from agentrun.utils.exception import ServerError from agentrun.utils.log import logger @@ -453,13 +454,79 @@ async def check_health_async(self): # Browser API Methods # ======================================== - def get_cdp_url(self, record: Optional[bool] = False): - """Get CDP WebSocket URL for browser automation.""" - return self.data_api.get_cdp_url(record=record) + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... + + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: + """Get CDP WebSocket URL for browser automation. + + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + config: Optional config override / 可选的配置覆盖 + """ + return self.data_api.get_cdp_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload] - def get_vnc_url(self, record: Optional[bool] = False): - """Get VNC WebSocket URL for live view.""" - return self.data_api.get_vnc_url(record=record) + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... + + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: + """Get VNC WebSocket URL for live view. + + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + config: Optional config override / 可选的配置覆盖 + """ + return self.data_api.get_vnc_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload] def sync_playwright(self, record: Optional[bool] = False): """Get synchronous Playwright browser instance.""" diff --git a/agentrun/sandbox/__browser_sandbox_async_template.py b/agentrun/sandbox/__browser_sandbox_async_template.py index 374e0c5..7292110 100644 --- a/agentrun/sandbox/__browser_sandbox_async_template.py +++ b/agentrun/sandbox/__browser_sandbox_async_template.py @@ -6,10 +6,11 @@ import asyncio import time # noqa: F401 -from typing import Optional +from typing import Dict, Literal, Optional, overload, Tuple, Union from agentrun.sandbox.api import BrowserDataAPI from agentrun.sandbox.model import TemplateType +from agentrun.utils.config import Config from agentrun.utils.log import logger from .sandbox import Sandbox @@ -78,11 +79,79 @@ def data_api(self): async def check_health_async(self): return await self.data_api.check_health_async() - def get_cdp_url(self, record: Optional[bool] = False): - return self.data_api.get_cdp_url(record=record) + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... + + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: + """Get CDP WebSocket URL for browser automation. - def get_vnc_url(self, record: Optional[bool] = False): - return self.data_api.get_vnc_url(record=record) + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + config: Optional config override / 可选的配置覆盖 + """ + return self.data_api.get_cdp_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload] + + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... + + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: + """Get VNC WebSocket URL for live view. + + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + config: Optional config override / 可选的配置覆盖 + """ + return self.data_api.get_vnc_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload] def sync_playwright(self, record: Optional[bool] = False): return self.data_api.sync_playwright(record=record) diff --git a/agentrun/sandbox/aio_sandbox.py b/agentrun/sandbox/aio_sandbox.py index e39211e..5f2a83f 100644 --- a/agentrun/sandbox/aio_sandbox.py +++ b/agentrun/sandbox/aio_sandbox.py @@ -17,10 +17,11 @@ import asyncio import time # noqa: F401 -from typing import Optional +from typing import Dict, Literal, Optional, overload, Tuple, Union from agentrun.sandbox.api.aio_data import AioDataAPI from agentrun.sandbox.model import CodeLanguage, TemplateType +from agentrun.utils.config import Config from agentrun.utils.exception import ServerError from agentrun.utils.log import logger @@ -817,13 +818,79 @@ def check_health(self): # Browser API Methods # ======================================== - def get_cdp_url(self, record: Optional[bool] = False): - """Get CDP WebSocket URL for browser automation.""" - return self.data_api.get_cdp_url(record=record) + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... + + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: + """Get CDP WebSocket URL for browser automation. + + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + config: Optional config override / 可选的配置覆盖 + """ + return self.data_api.get_cdp_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload] - def get_vnc_url(self, record: Optional[bool] = False): - """Get VNC WebSocket URL for live view.""" - return self.data_api.get_vnc_url(record=record) + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... + + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: + """Get VNC WebSocket URL for live view. + + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + config: Optional config override / 可选的配置覆盖 + """ + return self.data_api.get_vnc_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload] def sync_playwright(self, record: Optional[bool] = False): """Get synchronous Playwright browser instance.""" diff --git a/agentrun/sandbox/api/__aio_data_async_template.py b/agentrun/sandbox/api/__aio_data_async_template.py index 129ee38..ca293aa 100644 --- a/agentrun/sandbox/api/__aio_data_async_template.py +++ b/agentrun/sandbox/api/__aio_data_async_template.py @@ -5,7 +5,7 @@ combining browser and code interpreter capabilities. """ -from typing import Any, Dict, Optional +from typing import Any, Dict, Literal, Optional, overload, Tuple, Union from urllib.parse import parse_qs, urlencode, urlparse from agentrun.sandbox.model import CodeLanguage @@ -36,63 +36,151 @@ def __init__( # Browser API Methods # ======================================== - def get_cdp_url(self, record: Optional[bool] = False): + def _assemble_ws_url( + self, base: str, ws_path: str, record: Optional[bool] = False + ) -> str: + path = ws_path.lstrip("/") + raw = "/".join( + part.strip("/") for part in [base, self.namespace, path] if part + ) + ws_url = raw.replace("http", "ws") + u = urlparse(ws_url) + query_dict = parse_qs(u.query) + query_dict["tenantId"] = [self.config.get_account_id()] + if record: + query_dict["recording"] = ["true"] + new_query = urlencode(query_dict, doseq=True) + return u._replace(query=new_query).geturl() + + def _build_ws_url( + self, ws_path: str, record: Optional[bool] = False + ) -> str: + return self._assemble_ws_url( + self.config.get_data_endpoint(), ws_path, record + ) + + def _build_ws_url_with_headers( + self, + ws_path: str, + record: Optional[bool] = False, + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + cfg = Config.with_configs(self.config, config) + url = self._assemble_ws_url(self.get_base_url(cfg), ws_path, record) + url, headers, _ = self.auth( + url=url, headers=cfg.get_headers(), config=cfg + ) + return url, headers + + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... + + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: """ Generate the WebSocket URL for Chrome DevTools Protocol (CDP) connection. + 生成 Chrome DevTools Protocol (CDP) 连接的 WebSocket URL。 - This method constructs a WebSocket URL by: - 1. Converting the HTTP endpoint to WebSocket protocol (ws://) - 2. Parsing the existing URL and query parameters - 3. Adding the session ID to the query parameters - 4. Reconstructing the complete WebSocket URL + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + If False (default), return only the URL string for backward compatibility. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + 当为 False（默认）时，仅返回 URL 字符串以保持向后兼容。 + config: Optional config override / 可选的配置覆盖 Returns: - str: The complete WebSocket URL for CDP automation connection, - including the session ID in the query parameters. + str or Tuple[str, Dict[str, str]]: CDP WebSocket URL, or (url, headers) tuple + when with_headers=True. Example: >>> api = AioDataAPI("sandbox123") >>> api.get_cdp_url() - 'ws://example.com/ws/automation?sessionId=session456' + 'wss://example.com/sandboxes/sandbox123/ws/automation?tenantId=123' + >>> url, headers = api.get_cdp_url(with_headers=True) """ - cdp_url = self.with_path("/ws/automation").replace("http", "ws") - u = urlparse(cdp_url) - query_dict = parse_qs(u.query) - query_dict["tenantId"] = [self.config.get_account_id()] - if record: - query_dict["recording"] = ["true"] - new_query = urlencode(query_dict, doseq=True) - new_u = u._replace(query=new_query) - return new_u.geturl() + if with_headers: + return self._build_ws_url_with_headers( + "/ws/automation", record=record, config=config + ) + return self._build_ws_url("/ws/automation", record=record) + + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... - def get_vnc_url(self, record: Optional[bool] = False): + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: """ Generate the WebSocket URL for VNC (Virtual Network Computing) live view connection. + 生成 VNC 实时预览连接的 WebSocket URL。 - This method constructs a WebSocket URL for real-time browser viewing by: - 1. Converting the HTTP endpoint to WebSocket protocol (ws://) - 2. Parsing the existing URL and query parameters - 3. Adding the session ID to the query parameters - 4. Reconstructing the complete WebSocket URL + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + If False (default), return only the URL string for backward compatibility. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + 当为 False（默认）时，仅返回 URL 字符串以保持向后兼容。 + config: Optional config override / 可选的配置覆盖 Returns: - str: The complete WebSocket URL for VNC live view connection, - including the session ID in the query parameters. + str or Tuple[str, Dict[str, str]]: VNC WebSocket URL, or (url, headers) tuple + when with_headers=True. Example: >>> api = AioDataAPI("sandbox123") >>> api.get_vnc_url() - 'ws://example.com/ws/liveview?sessionId=session456' + 'wss://example.com/sandboxes/sandbox123/ws/liveview?tenantId=123' + >>> url, headers = api.get_vnc_url(with_headers=True) """ - vnc_url = self.with_path("/ws/liveview").replace("http", "ws") - u = urlparse(vnc_url) - query_dict = parse_qs(u.query) - query_dict["tenantId"] = [self.config.get_account_id()] - if record: - query_dict["recording"] = ["true"] - new_query = urlencode(query_dict, doseq=True) - new_u = u._replace(query=new_query) - return new_u.geturl() + if with_headers: + return self._build_ws_url_with_headers( + "/ws/liveview", record=record, config=config + ) + return self._build_ws_url("/ws/liveview", record=record) def sync_playwright( self, @@ -102,11 +190,8 @@ def sync_playwright( ): from .playwright_sync import BrowserPlaywrightSync - cfg = Config.with_configs(self.config, config) - - url = self.get_cdp_url(record=record) - url, headers, _ = self.auth( - url=url, headers=cfg.get_headers(), config=cfg + url, headers = self._build_ws_url_with_headers( + "/ws/automation", record=record, config=config ) return BrowserPlaywrightSync( url, @@ -122,11 +207,8 @@ def async_playwright( ): from .playwright_async import BrowserPlaywrightAsync - cfg = Config.with_configs(self.config, config) - - url = self.get_cdp_url(record=record) - url, headers, _ = self.auth( - url=url, headers=cfg.get_headers(), config=cfg + url, headers = self._build_ws_url_with_headers( + "/ws/automation", record=record, config=config ) return BrowserPlaywrightAsync( url, diff --git a/agentrun/sandbox/api/__browser_data_async_template.py b/agentrun/sandbox/api/__browser_data_async_template.py index cf215b5..957038b 100644 --- a/agentrun/sandbox/api/__browser_data_async_template.py +++ b/agentrun/sandbox/api/__browser_data_async_template.py @@ -4,7 +4,7 @@ This template is used to generate browser sandbox data API code. """ -from typing import Optional +from typing import Dict, Literal, Optional, overload, Tuple, Union from urllib.parse import parse_qs, urlencode, urlparse from agentrun.utils.config import Config @@ -25,63 +25,151 @@ def __init__( config=config, ) - def get_cdp_url(self, record: Optional[bool] = False): + def _assemble_ws_url( + self, base: str, ws_path: str, record: Optional[bool] = False + ) -> str: + path = ws_path.lstrip("/") + raw = "/".join( + part.strip("/") for part in [base, self.namespace, path] if part + ) + ws_url = raw.replace("http", "ws") + u = urlparse(ws_url) + query_dict = parse_qs(u.query) + query_dict["tenantId"] = [self.config.get_account_id()] + if record: + query_dict["recording"] = ["true"] + new_query = urlencode(query_dict, doseq=True) + return u._replace(query=new_query).geturl() + + def _build_ws_url( + self, ws_path: str, record: Optional[bool] = False + ) -> str: + return self._assemble_ws_url( + self.config.get_data_endpoint(), ws_path, record + ) + + def _build_ws_url_with_headers( + self, + ws_path: str, + record: Optional[bool] = False, + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + cfg = Config.with_configs(self.config, config) + url = self._assemble_ws_url(self.get_base_url(cfg), ws_path, record) + url, headers, _ = self.auth( + url=url, headers=cfg.get_headers(), config=cfg + ) + return url, headers + + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... + + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: """ Generate the WebSocket URL for Chrome DevTools Protocol (CDP) connection. + 生成 Chrome DevTools Protocol (CDP) 连接的 WebSocket URL。 - This method constructs a WebSocket URL by: - 1. Converting the HTTP endpoint to WebSocket protocol (ws://) - 2. Parsing the existing URL and query parameters - 3. Adding the session ID to the query parameters - 4. Reconstructing the complete WebSocket URL + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + If False (default), return only the URL string for backward compatibility. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + 当为 False（默认）时，仅返回 URL 字符串以保持向后兼容。 + config: Optional config override / 可选的配置覆盖 Returns: - str: The complete WebSocket URL for CDP automation connection, - including the session ID in the query parameters. + str or Tuple[str, Dict[str, str]]: CDP WebSocket URL, or (url, headers) tuple + when with_headers=True. Example: - >>> api = BrowserDataAPI("browser123", "session456") + >>> api = BrowserDataAPI("browser123") >>> api.get_cdp_url() - 'ws://example.com/ws/automation?sessionId=session456' + 'wss://example.com/sandboxes/browser123/ws/automation?tenantId=123' + >>> url, headers = api.get_cdp_url(with_headers=True) """ - cdp_url = self.with_path("/ws/automation").replace("http", "ws") - u = urlparse(cdp_url) - query_dict = parse_qs(u.query) - query_dict["tenantId"] = [self.config.get_account_id()] - if record: - query_dict["recording"] = ["true"] - new_query = urlencode(query_dict, doseq=True) - new_u = u._replace(query=new_query) - return new_u.geturl() + if with_headers: + return self._build_ws_url_with_headers( + "/ws/automation", record=record, config=config + ) + return self._build_ws_url("/ws/automation", record=record) + + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... - def get_vnc_url(self, record: Optional[bool] = False): + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: """ Generate the WebSocket URL for VNC (Virtual Network Computing) live view connection. + 生成 VNC 实时预览连接的 WebSocket URL。 - This method constructs a WebSocket URL for real-time browser viewing by: - 1. Converting the HTTP endpoint to WebSocket protocol (ws://) - 2. Parsing the existing URL and query parameters - 3. Adding the session ID to the query parameters - 4. Reconstructing the complete WebSocket URL + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + If False (default), return only the URL string for backward compatibility. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + 当为 False（默认）时，仅返回 URL 字符串以保持向后兼容。 + config: Optional config override / 可选的配置覆盖 Returns: - str: The complete WebSocket URL for VNC live view connection, - including the session ID in the query parameters. + str or Tuple[str, Dict[str, str]]: VNC WebSocket URL, or (url, headers) tuple + when with_headers=True. Example: - >>> api = BrowserDataAPI("browser123", "session456") + >>> api = BrowserDataAPI("browser123") >>> api.get_vnc_url() - 'ws://example.com/ws/liveview?sessionId=session456' + 'wss://example.com/sandboxes/browser123/ws/liveview?tenantId=123' + >>> url, headers = api.get_vnc_url(with_headers=True) """ - vnc_url = self.with_path("/ws/liveview").replace("http", "ws") - u = urlparse(vnc_url) - query_dict = parse_qs(u.query) - query_dict["tenantId"] = [self.config.get_account_id()] - if record: - query_dict["recording"] = ["true"] - new_query = urlencode(query_dict, doseq=True) - new_u = u._replace(query=new_query) - return new_u.geturl() + if with_headers: + return self._build_ws_url_with_headers( + "/ws/liveview", record=record, config=config + ) + return self._build_ws_url("/ws/liveview", record=record) def sync_playwright( self, @@ -91,11 +179,8 @@ def sync_playwright( ): from .playwright_sync import BrowserPlaywrightSync - cfg = Config.with_configs(self.config, config) - - url = self.get_cdp_url(record=record) - url, headers, _ = self.auth( - url=url, headers=cfg.get_headers(), config=cfg + url, headers = self._build_ws_url_with_headers( + "/ws/automation", record=record, config=config ) return BrowserPlaywrightSync( url, @@ -111,11 +196,8 @@ def async_playwright( ): from .playwright_async import BrowserPlaywrightAsync - cfg = Config.with_configs(self.config, config) - - url = self.get_cdp_url(record=record) - url, headers, _ = self.auth( - url=url, headers=cfg.get_headers(), config=cfg + url, headers = self._build_ws_url_with_headers( + "/ws/automation", record=record, config=config ) return BrowserPlaywrightAsync( url, diff --git a/agentrun/sandbox/api/aio_data.py b/agentrun/sandbox/api/aio_data.py index d271db4..18a1c48 100644 --- a/agentrun/sandbox/api/aio_data.py +++ b/agentrun/sandbox/api/aio_data.py @@ -15,7 +15,7 @@ combining browser and code interpreter capabilities. """ -from typing import Any, Dict, Optional +from typing import Any, Dict, Literal, Optional, overload, Tuple, Union from urllib.parse import parse_qs, urlencode, urlparse from agentrun.sandbox.model import CodeLanguage @@ -46,63 +46,151 @@ def __init__( # Browser API Methods # ======================================== - def get_cdp_url(self, record: Optional[bool] = False): + def _assemble_ws_url( + self, base: str, ws_path: str, record: Optional[bool] = False + ) -> str: + path = ws_path.lstrip("/") + raw = "/".join( + part.strip("/") for part in [base, self.namespace, path] if part + ) + ws_url = raw.replace("http", "ws") + u = urlparse(ws_url) + query_dict = parse_qs(u.query) + query_dict["tenantId"] = [self.config.get_account_id()] + if record: + query_dict["recording"] = ["true"] + new_query = urlencode(query_dict, doseq=True) + return u._replace(query=new_query).geturl() + + def _build_ws_url( + self, ws_path: str, record: Optional[bool] = False + ) -> str: + return self._assemble_ws_url( + self.config.get_data_endpoint(), ws_path, record + ) + + def _build_ws_url_with_headers( + self, + ws_path: str, + record: Optional[bool] = False, + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + cfg = Config.with_configs(self.config, config) + url = self._assemble_ws_url(self.get_base_url(cfg), ws_path, record) + url, headers, _ = self.auth( + url=url, headers=cfg.get_headers(), config=cfg + ) + return url, headers + + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... + + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: """ Generate the WebSocket URL for Chrome DevTools Protocol (CDP) connection. + 生成 Chrome DevTools Protocol (CDP) 连接的 WebSocket URL。 - This method constructs a WebSocket URL by: - 1. Converting the HTTP endpoint to WebSocket protocol (ws://) - 2. Parsing the existing URL and query parameters - 3. Adding the session ID to the query parameters - 4. Reconstructing the complete WebSocket URL + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + If False (default), return only the URL string for backward compatibility. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + 当为 False（默认）时，仅返回 URL 字符串以保持向后兼容。 + config: Optional config override / 可选的配置覆盖 Returns: - str: The complete WebSocket URL for CDP automation connection, - including the session ID in the query parameters. + str or Tuple[str, Dict[str, str]]: CDP WebSocket URL, or (url, headers) tuple + when with_headers=True. Example: >>> api = AioDataAPI("sandbox123") >>> api.get_cdp_url() - 'ws://example.com/ws/automation?sessionId=session456' + 'wss://example.com/sandboxes/sandbox123/ws/automation?tenantId=123' + >>> url, headers = api.get_cdp_url(with_headers=True) """ - cdp_url = self.with_path("/ws/automation").replace("http", "ws") - u = urlparse(cdp_url) - query_dict = parse_qs(u.query) - query_dict["tenantId"] = [self.config.get_account_id()] - if record: - query_dict["recording"] = ["true"] - new_query = urlencode(query_dict, doseq=True) - new_u = u._replace(query=new_query) - return new_u.geturl() + if with_headers: + return self._build_ws_url_with_headers( + "/ws/automation", record=record, config=config + ) + return self._build_ws_url("/ws/automation", record=record) + + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... - def get_vnc_url(self, record: Optional[bool] = False): + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: """ Generate the WebSocket URL for VNC (Virtual Network Computing) live view connection. + 生成 VNC 实时预览连接的 WebSocket URL。 - This method constructs a WebSocket URL for real-time browser viewing by: - 1. Converting the HTTP endpoint to WebSocket protocol (ws://) - 2. Parsing the existing URL and query parameters - 3. Adding the session ID to the query parameters - 4. Reconstructing the complete WebSocket URL + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + If False (default), return only the URL string for backward compatibility. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + 当为 False（默认）时，仅返回 URL 字符串以保持向后兼容。 + config: Optional config override / 可选的配置覆盖 Returns: - str: The complete WebSocket URL for VNC live view connection, - including the session ID in the query parameters. + str or Tuple[str, Dict[str, str]]: VNC WebSocket URL, or (url, headers) tuple + when with_headers=True. Example: >>> api = AioDataAPI("sandbox123") >>> api.get_vnc_url() - 'ws://example.com/ws/liveview?sessionId=session456' + 'wss://example.com/sandboxes/sandbox123/ws/liveview?tenantId=123' + >>> url, headers = api.get_vnc_url(with_headers=True) """ - vnc_url = self.with_path("/ws/liveview").replace("http", "ws") - u = urlparse(vnc_url) - query_dict = parse_qs(u.query) - query_dict["tenantId"] = [self.config.get_account_id()] - if record: - query_dict["recording"] = ["true"] - new_query = urlencode(query_dict, doseq=True) - new_u = u._replace(query=new_query) - return new_u.geturl() + if with_headers: + return self._build_ws_url_with_headers( + "/ws/liveview", record=record, config=config + ) + return self._build_ws_url("/ws/liveview", record=record) def sync_playwright( self, @@ -112,11 +200,8 @@ def sync_playwright( ): from .playwright_sync import BrowserPlaywrightSync - cfg = Config.with_configs(self.config, config) - - url = self.get_cdp_url(record=record) - url, headers, _ = self.auth( - url=url, headers=cfg.get_headers(), config=cfg + url, headers = self._build_ws_url_with_headers( + "/ws/automation", record=record, config=config ) return BrowserPlaywrightSync( url, @@ -132,11 +217,8 @@ def async_playwright( ): from .playwright_async import BrowserPlaywrightAsync - cfg = Config.with_configs(self.config, config) - - url = self.get_cdp_url(record=record) - url, headers, _ = self.auth( - url=url, headers=cfg.get_headers(), config=cfg + url, headers = self._build_ws_url_with_headers( + "/ws/automation", record=record, config=config ) return BrowserPlaywrightAsync( url, diff --git a/agentrun/sandbox/api/browser_data.py b/agentrun/sandbox/api/browser_data.py index 2523605..e40aa99 100644 --- a/agentrun/sandbox/api/browser_data.py +++ b/agentrun/sandbox/api/browser_data.py @@ -14,7 +14,7 @@ This template is used to generate browser sandbox data API code. """ -from typing import Optional +from typing import Dict, Literal, Optional, overload, Tuple, Union from urllib.parse import parse_qs, urlencode, urlparse from agentrun.utils.config import Config @@ -35,63 +35,151 @@ def __init__( config=config, ) - def get_cdp_url(self, record: Optional[bool] = False): + def _assemble_ws_url( + self, base: str, ws_path: str, record: Optional[bool] = False + ) -> str: + path = ws_path.lstrip("/") + raw = "/".join( + part.strip("/") for part in [base, self.namespace, path] if part + ) + ws_url = raw.replace("http", "ws") + u = urlparse(ws_url) + query_dict = parse_qs(u.query) + query_dict["tenantId"] = [self.config.get_account_id()] + if record: + query_dict["recording"] = ["true"] + new_query = urlencode(query_dict, doseq=True) + return u._replace(query=new_query).geturl() + + def _build_ws_url( + self, ws_path: str, record: Optional[bool] = False + ) -> str: + return self._assemble_ws_url( + self.config.get_data_endpoint(), ws_path, record + ) + + def _build_ws_url_with_headers( + self, + ws_path: str, + record: Optional[bool] = False, + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + cfg = Config.with_configs(self.config, config) + url = self._assemble_ws_url(self.get_base_url(cfg), ws_path, record) + url, headers, _ = self.auth( + url=url, headers=cfg.get_headers(), config=cfg + ) + return url, headers + + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... + + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: """ Generate the WebSocket URL for Chrome DevTools Protocol (CDP) connection. + 生成 Chrome DevTools Protocol (CDP) 连接的 WebSocket URL。 - This method constructs a WebSocket URL by: - 1. Converting the HTTP endpoint to WebSocket protocol (ws://) - 2. Parsing the existing URL and query parameters - 3. Adding the session ID to the query parameters - 4. Reconstructing the complete WebSocket URL + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + If False (default), return only the URL string for backward compatibility. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + 当为 False（默认）时，仅返回 URL 字符串以保持向后兼容。 + config: Optional config override / 可选的配置覆盖 Returns: - str: The complete WebSocket URL for CDP automation connection, - including the session ID in the query parameters. + str or Tuple[str, Dict[str, str]]: CDP WebSocket URL, or (url, headers) tuple + when with_headers=True. Example: - >>> api = BrowserDataAPI("browser123", "session456") + >>> api = BrowserDataAPI("browser123") >>> api.get_cdp_url() - 'ws://example.com/ws/automation?sessionId=session456' + 'wss://example.com/sandboxes/browser123/ws/automation?tenantId=123' + >>> url, headers = api.get_cdp_url(with_headers=True) """ - cdp_url = self.with_path("/ws/automation").replace("http", "ws") - u = urlparse(cdp_url) - query_dict = parse_qs(u.query) - query_dict["tenantId"] = [self.config.get_account_id()] - if record: - query_dict["recording"] = ["true"] - new_query = urlencode(query_dict, doseq=True) - new_u = u._replace(query=new_query) - return new_u.geturl() + if with_headers: + return self._build_ws_url_with_headers( + "/ws/automation", record=record, config=config + ) + return self._build_ws_url("/ws/automation", record=record) + + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... - def get_vnc_url(self, record: Optional[bool] = False): + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: """ Generate the WebSocket URL for VNC (Virtual Network Computing) live view connection. + 生成 VNC 实时预览连接的 WebSocket URL。 - This method constructs a WebSocket URL for real-time browser viewing by: - 1. Converting the HTTP endpoint to WebSocket protocol (ws://) - 2. Parsing the existing URL and query parameters - 3. Adding the session ID to the query parameters - 4. Reconstructing the complete WebSocket URL + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + If False (default), return only the URL string for backward compatibility. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + 当为 False（默认）时，仅返回 URL 字符串以保持向后兼容。 + config: Optional config override / 可选的配置覆盖 Returns: - str: The complete WebSocket URL for VNC live view connection, - including the session ID in the query parameters. + str or Tuple[str, Dict[str, str]]: VNC WebSocket URL, or (url, headers) tuple + when with_headers=True. Example: - >>> api = BrowserDataAPI("browser123", "session456") + >>> api = BrowserDataAPI("browser123") >>> api.get_vnc_url() - 'ws://example.com/ws/liveview?sessionId=session456' + 'wss://example.com/sandboxes/browser123/ws/liveview?tenantId=123' + >>> url, headers = api.get_vnc_url(with_headers=True) """ - vnc_url = self.with_path("/ws/liveview").replace("http", "ws") - u = urlparse(vnc_url) - query_dict = parse_qs(u.query) - query_dict["tenantId"] = [self.config.get_account_id()] - if record: - query_dict["recording"] = ["true"] - new_query = urlencode(query_dict, doseq=True) - new_u = u._replace(query=new_query) - return new_u.geturl() + if with_headers: + return self._build_ws_url_with_headers( + "/ws/liveview", record=record, config=config + ) + return self._build_ws_url("/ws/liveview", record=record) def sync_playwright( self, @@ -101,11 +189,8 @@ def sync_playwright( ): from .playwright_sync import BrowserPlaywrightSync - cfg = Config.with_configs(self.config, config) - - url = self.get_cdp_url(record=record) - url, headers, _ = self.auth( - url=url, headers=cfg.get_headers(), config=cfg + url, headers = self._build_ws_url_with_headers( + "/ws/automation", record=record, config=config ) return BrowserPlaywrightSync( url, @@ -121,11 +206,8 @@ def async_playwright( ): from .playwright_async import BrowserPlaywrightAsync - cfg = Config.with_configs(self.config, config) - - url = self.get_cdp_url(record=record) - url, headers, _ = self.auth( - url=url, headers=cfg.get_headers(), config=cfg + url, headers = self._build_ws_url_with_headers( + "/ws/automation", record=record, config=config ) return BrowserPlaywrightAsync( url, diff --git a/agentrun/sandbox/browser_sandbox.py b/agentrun/sandbox/browser_sandbox.py index e43b801..55b3b8a 100644 --- a/agentrun/sandbox/browser_sandbox.py +++ b/agentrun/sandbox/browser_sandbox.py @@ -16,10 +16,11 @@ import asyncio import time # noqa: F401 -from typing import Optional +from typing import Dict, Literal, Optional, overload, Tuple, Union from agentrun.sandbox.api import BrowserDataAPI from agentrun.sandbox.model import TemplateType +from agentrun.utils.config import Config from agentrun.utils.log import logger from .sandbox import Sandbox @@ -134,11 +135,79 @@ async def check_health_async(self): def check_health(self): return self.data_api.check_health() - def get_cdp_url(self, record: Optional[bool] = False): - return self.data_api.get_cdp_url(record=record) + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... + + @overload + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_cdp_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: + """Get CDP WebSocket URL for browser automation. - def get_vnc_url(self, record: Optional[bool] = False): - return self.data_api.get_vnc_url(record=record) + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + config: Optional config override / 可选的配置覆盖 + """ + return self.data_api.get_cdp_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload] + + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[True], + config: Optional[Config] = None, + ) -> Tuple[str, Dict[str, str]]: + ... + + @overload + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: Literal[False] = False, + config: Optional[Config] = None, + ) -> str: + ... + + def get_vnc_url( + self, + record: Optional[bool] = False, + *, + with_headers: bool = False, + config: Optional[Config] = None, + ) -> Union[str, Tuple[str, Dict[str, str]]]: + """Get VNC WebSocket URL for live view. + + Args: + record: Whether to enable recording / 是否启用录制 + with_headers: If True, return (url, headers) tuple with authentication headers. + 当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。 + config: Optional config override / 可选的配置覆盖 + """ + return self.data_api.get_vnc_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload] def sync_playwright(self, record: Optional[bool] = False): return self.data_api.sync_playwright(record=record) diff --git a/agentrun/sandbox/custom_sandbox.py b/agentrun/sandbox/custom_sandbox.py index 5e8eadd..d232206 100644 --- a/agentrun/sandbox/custom_sandbox.py +++ b/agentrun/sandbox/custom_sandbox.py @@ -2,7 +2,6 @@ from agentrun.sandbox.model import TemplateType from agentrun.utils.config import Config -from agentrun.utils.data_api import DataAPI, ResourceType from .sandbox import Sandbox @@ -13,12 +12,12 @@ class CustomSandbox(Sandbox): _template_type = TemplateType.CUSTOM def get_base_url(self, config: Optional[Config] = None): - """Get the base URL for the custom sandbox template.""" - api = DataAPI( - resource_name="", - resource_type=ResourceType.Template, - namespace="sandboxes", - config=config, - ) + """Get the base URL for the custom sandbox template. - return api.with_path("") + Returns the non-RAM data endpoint so that the URL can be used + directly without requiring RAM signature headers. + 返回非 RAM 的数据端点，以便用户可以直接使用该 URL 而无需附带 RAM 签名头。 + """ + cfg = Config.with_configs(config) + base = cfg.get_data_endpoint() + return "/".join(part.strip("/") for part in [base, "sandboxes"] if part) diff --git a/tests/unittests/sandbox/api/test_aio_data.py b/tests/unittests/sandbox/api/test_aio_data.py index 1922134..2f74ba8 100644 --- a/tests/unittests/sandbox/api/test_aio_data.py +++ b/tests/unittests/sandbox/api/test_aio_data.py @@ -6,6 +6,10 @@ from agentrun.sandbox.api.aio_data import AioDataAPI from agentrun.sandbox.model import CodeLanguage +from agentrun.utils.config import Config + +_DATA_ENDPOINT = "https://account123.agentrun-data.cn-hangzhou.aliyuncs.com" +_RAM_ENDPOINT = "https://account123-ram.agentrun-data.cn-hangzhou.aliyuncs.com" @pytest.fixture @@ -13,17 +17,27 @@ def api(): with patch.object(AioDataAPI, "__init__", lambda self, **kw: None): obj = AioDataAPI.__new__(AioDataAPI) obj.sandbox_id = "sb-aio-1" - obj.config = MagicMock() - obj.config.get_account_id.return_value = "account123" + obj.config = Config( + account_id="account123", + data_endpoint=_DATA_ENDPOINT, + ) obj.access_token = "tok" obj.access_token_map = {} obj.resource_name = "sb-aio-1" + obj.namespace = "sandboxes/sb-aio-1" obj.with_path = MagicMock( - side_effect=lambda p: f"http://host.com/ns{p}?sig=abc" - ) - obj.auth = MagicMock( - return_value=("tok", {"Authorization": "Bearer tok"}, None) + side_effect=lambda p, **kw: f"http://host.com/ns{p}?sig=abc" ) + obj.get_base_url = MagicMock(return_value=_RAM_ENDPOINT) + + def _auth_side_effect(url="", headers=None, query=None, **kw): + return ( + url, + {"Authorization": "Bearer tok", **(headers or {})}, + query, + ) + + obj.auth = MagicMock(side_effect=_auth_side_effect) obj.get = MagicMock(return_value={"ok": True}) obj.get_async = AsyncMock(return_value={"ok": True}) obj.post = MagicMock(return_value={"ok": True}) @@ -62,28 +76,40 @@ class TestAioCdpUrl: def test_get_cdp_url_no_record(self, api): url = api.get_cdp_url() - api.with_path.assert_called_once_with("/ws/automation") - assert "ws://" in url + assert "wss://" in url assert "tenantId=account123" in url assert "recording" not in url + assert "ws/automation" in url + assert "-ram" not in url def test_get_cdp_url_with_record(self, api): url = api.get_cdp_url(record=True) assert "recording=true" in url + def test_get_cdp_url_with_headers(self, api): + url, headers = api.get_cdp_url(with_headers=True) + assert "Authorization" in headers + assert "ws/automation" in url + class TestAioVncUrl: def test_get_vnc_url_no_record(self, api): url = api.get_vnc_url() - api.with_path.assert_called_once_with("/ws/liveview") - assert "ws://" in url + assert "wss://" in url assert "tenantId=account123" in url + assert "ws/liveview" in url + assert "-ram" not in url def test_get_vnc_url_with_record(self, api): url = api.get_vnc_url(record=True) assert "recording=true" in url + def test_get_vnc_url_with_headers(self, api): + url, headers = api.get_vnc_url(with_headers=True) + assert "Authorization" in headers + assert "ws/liveview" in url + class TestAioPlaywright: diff --git a/tests/unittests/sandbox/api/test_browser_data.py b/tests/unittests/sandbox/api/test_browser_data.py index 354ffd1..f2f52a8 100644 --- a/tests/unittests/sandbox/api/test_browser_data.py +++ b/tests/unittests/sandbox/api/test_browser_data.py @@ -5,6 +5,10 @@ import pytest from agentrun.sandbox.api.browser_data import BrowserDataAPI +from agentrun.utils.config import Config + +_DATA_ENDPOINT = "https://account123.agentrun-data.cn-hangzhou.aliyuncs.com" +_RAM_ENDPOINT = "https://account123-ram.agentrun-data.cn-hangzhou.aliyuncs.com" @pytest.fixture @@ -12,17 +16,28 @@ def api(): with patch.object(BrowserDataAPI, "__init__", lambda self, **kw: None): obj = BrowserDataAPI.__new__(BrowserDataAPI) obj.sandbox_id = "sb-1" - obj.config = MagicMock() - obj.config.get_account_id.return_value = "account123" + obj.config = Config( + account_id="account123", + data_endpoint=_DATA_ENDPOINT, + ) obj.access_token = "tok" obj.access_token_map = {} obj.resource_name = "sb-1" + obj.namespace = "sandboxes/sb-1" obj.with_path = MagicMock( - side_effect=lambda p: f"http://host.com/ns{p}?sig=abc" - ) - obj.auth = MagicMock( - return_value=("tok", {"Authorization": "Bearer tok"}, None) + side_effect=lambda p, **kw: f"http://host.com/ns{p}?sig=abc" ) + obj.get_base_url = MagicMock(return_value=_RAM_ENDPOINT) + + def _auth_side_effect(url="", headers=None, query=None, **kw): + return ( + url, + {"Authorization": "Bearer tok", **(headers or {})}, + query, + ) + + obj.auth = MagicMock(side_effect=_auth_side_effect) + obj.get = MagicMock(return_value=[]) obj.get_async = AsyncMock(return_value=[]) obj.delete = MagicMock(return_value={"ok": True}) @@ -50,29 +65,41 @@ class TestCdpUrl: def test_get_cdp_url_no_record(self, api): url = api.get_cdp_url() - api.with_path.assert_called_once_with("/ws/automation") - assert "ws://" in url + assert "wss://" in url assert "tenantId=account123" in url assert "recording" not in url + assert "ws/automation" in url + assert "-ram" not in url def test_get_cdp_url_with_record(self, api): url = api.get_cdp_url(record=True) assert "recording=true" in url assert "tenantId=account123" in url + def test_get_cdp_url_with_headers(self, api): + url, headers = api.get_cdp_url(with_headers=True) + assert "Authorization" in headers + assert "ws/automation" in url + class TestVncUrl: def test_get_vnc_url_no_record(self, api): url = api.get_vnc_url() - api.with_path.assert_called_once_with("/ws/liveview") - assert "ws://" in url + assert "wss://" in url assert "tenantId=account123" in url + assert "ws/liveview" in url + assert "-ram" not in url def test_get_vnc_url_with_record(self, api): url = api.get_vnc_url(record=True) assert "recording=true" in url + def test_get_vnc_url_with_headers(self, api): + url, headers = api.get_vnc_url(with_headers=True) + assert "Authorization" in headers + assert "ws/liveview" in url + class TestPlaywright: diff --git a/tests/unittests/sandbox/test_browser_sandbox.py b/tests/unittests/sandbox/test_browser_sandbox.py index 55efa8b..1e7e540 100644 --- a/tests/unittests/sandbox/test_browser_sandbox.py +++ b/tests/unittests/sandbox/test_browser_sandbox.py @@ -56,7 +56,9 @@ def test_get_cdp_url(self): sb = _make_sandbox() sb.data_api.get_cdp_url.return_value = "ws://example.com/ws/automation" assert sb.get_cdp_url(record=True) == "ws://example.com/ws/automation" - sb.data_api.get_cdp_url.assert_called_once_with(record=True) + sb.data_api.get_cdp_url.assert_called_once_with( + record=True, with_headers=False, config=None + ) def test_get_vnc_url(self): sb = _make_sandbox() diff --git a/tests/unittests/sandbox/test_custom_sandbox.py b/tests/unittests/sandbox/test_custom_sandbox.py index b966c12..ceb86b6 100644 --- a/tests/unittests/sandbox/test_custom_sandbox.py +++ b/tests/unittests/sandbox/test_custom_sandbox.py @@ -1,9 +1,5 @@ """Tests for agentrun.sandbox.custom_sandbox module.""" -from unittest.mock import MagicMock, patch - -import pytest - from agentrun.sandbox.custom_sandbox import CustomSandbox from agentrun.sandbox.model import TemplateType @@ -16,24 +12,26 @@ def test_template_type(self): == TemplateType.CUSTOM ) - @patch("agentrun.sandbox.custom_sandbox.DataAPI") - def test_get_base_url(self, mock_data_api_cls): - mock_api = MagicMock() - mock_api.with_path.return_value = "https://example.com/sandboxes" - mock_data_api_cls.return_value = mock_api + def test_get_base_url(self): + from agentrun.utils.config import Config + cfg = Config( + data_endpoint=( + "https://account123.agentrun-data.cn-hangzhou.aliyuncs.com" + ) + ) sb = CustomSandbox.model_construct(sandbox_id="sb-1") - result = sb.get_base_url() - assert result == "https://example.com/sandboxes" - mock_api.with_path.assert_called_once_with("") + result = sb.get_base_url(config=cfg) + assert ( + result + == "https://account123.agentrun-data.cn-hangzhou.aliyuncs.com/sandboxes" + ) + assert "-ram" not in result - @patch("agentrun.sandbox.custom_sandbox.DataAPI") - def test_get_base_url_with_config(self, mock_data_api_cls): - mock_api = MagicMock() - mock_api.with_path.return_value = "https://custom.com" - mock_data_api_cls.return_value = mock_api + def test_get_base_url_with_config(self): + from agentrun.utils.config import Config + cfg = Config(data_endpoint="https://custom.com") sb = CustomSandbox.model_construct(sandbox_id="sb-1") - config = MagicMock() - result = sb.get_base_url(config=config) - assert result == "https://custom.com" + result = sb.get_base_url(config=cfg) + assert result == "https://custom.com/sandboxes" From 945368f7fda732233e791c20de0c967bb5d02bfb Mon Sep 17 00:00:00 2001 From: OhYee Date: Wed, 1 Apr 2026 18:11:36 +0800 Subject: [PATCH 2/4] Update agentrun/sandbox/api/browser_data.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: OhYee --- agentrun/sandbox/api/browser_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agentrun/sandbox/api/browser_data.py b/agentrun/sandbox/api/browser_data.py index e40aa99..08c6d70 100644 --- a/agentrun/sandbox/api/browser_data.py +++ b/agentrun/sandbox/api/browser_data.py @@ -179,7 +179,7 @@ def get_vnc_url( return self._build_ws_url_with_headers( "/ws/liveview", record=record, config=config ) - return self._build_ws_url("/ws/liveview", record=record) + return self._build_ws_url("/ws/liveview", record=record, config=config) def sync_playwright( self, From e62b7ccf0ee0b9d498b3ffc25711e320d5f20f9c Mon Sep 17 00:00:00 2001 From: OhYee Date: Wed, 1 Apr 2026 18:13:22 +0800 Subject: [PATCH 3/4] Apply suggestion from @Copilot Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: OhYee --- agentrun/sandbox/api/aio_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agentrun/sandbox/api/aio_data.py b/agentrun/sandbox/api/aio_data.py index 18a1c48..32438ea 100644 --- a/agentrun/sandbox/api/aio_data.py +++ b/agentrun/sandbox/api/aio_data.py @@ -135,7 +135,7 @@ def get_cdp_url( return self._build_ws_url_with_headers( "/ws/automation", record=record, config=config ) - return self._build_ws_url("/ws/automation", record=record) + return self._build_ws_url("/ws/automation", record=record, config=config) @overload def get_vnc_url( From afcdfb13d97aa67ef3ee0fd98679fcf4b7be9896 Mon Sep 17 00:00:00 2001 From: OhYee Date: Wed, 1 Apr 2026 18:14:13 +0800 Subject: [PATCH 4/4] Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: OhYee --- agentrun/sandbox/api/aio_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agentrun/sandbox/api/aio_data.py b/agentrun/sandbox/api/aio_data.py index 32438ea..8cf3717 100644 --- a/agentrun/sandbox/api/aio_data.py +++ b/agentrun/sandbox/api/aio_data.py @@ -190,7 +190,7 @@ def get_vnc_url( return self._build_ws_url_with_headers( "/ws/liveview", record=record, config=config ) - return self._build_ws_url("/ws/liveview", record=record) + return self._build_ws_url("/ws/liveview", record=record, config=config) def sync_playwright( self,