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..93994ce 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,153 @@ 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, + config: Optional[Config] = None, + ) -> str: + cfg = Config.with_configs(self.config, config) + return self._assemble_ws_url(cfg.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]]: + ... + + @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): + 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 +192,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 +209,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..26fe489 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,153 @@ 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, + config: Optional[Config] = None, + ) -> str: + cfg = Config.with_configs(self.config, config) + return self._assemble_ws_url(cfg.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 +181,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 +198,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..6581bc0 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,153 @@ 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, + config: Optional[Config] = None, + ) -> str: + cfg = Config.with_configs(self.config, config) + return self._assemble_ws_url(cfg.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]]: + ... + + @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): + 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 +202,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 +219,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..aa43e2f 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,153 @@ 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, + config: Optional[Config] = None, + ) -> str: + cfg = Config.with_configs(self.config, config) + return self._assemble_ws_url(cfg.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 +191,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 +208,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/agentrun/utils/ram_signature/signer.py b/agentrun/utils/ram_signature/signer.py index ce0b126..c91ddf8 100644 --- a/agentrun/utils/ram_signature/signer.py +++ b/agentrun/utils/ram_signature/signer.py @@ -12,7 +12,7 @@ import hashlib import hmac from typing import Optional -from urllib.parse import quote, urlparse +from urllib.parse import quote, unquote, urlparse ALGORITHM = "AGENTRUN4-HMAC-SHA256" UNSIGNED_PAYLOAD = "UNSIGNED-PAYLOAD" @@ -156,7 +156,7 @@ def get_agentrun_signed_headers( for pair in parsed.query.split("&"): if "=" in pair: k, v = pair.split("=", 1) - query_params[k] = v + query_params[unquote(k)] = unquote(v) now = sign_time if sign_time is not None else datetime.now(timezone.utc) if now.tzinfo is None: 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"