Skip to content

Commit 3c54da1

Browse files
OhYeeOhYee
authored
Merge pull request #78 from Serverless-Devs/fix-new-auth
feat(browser-api): enhance WebSocket URL generation with header support
2 parents 13944fb + b49ff60 commit 3c54da1

14 files changed

+928
-268
lines changed

‎agentrun/sandbox/__aio_sandbox_async_template.py‎

Lines changed: 74 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,11 @@
77

88
import asyncio
99
import time # noqa: F401
10-
from typing import Optional
10+
from typing import Dict, Literal, Optional, overload, Tuple, Union
1111

1212
from agentrun.sandbox.api.aio_data import AioDataAPI
1313
from agentrun.sandbox.model import CodeLanguage, TemplateType
14+
from agentrun.utils.config import Config
1415
from agentrun.utils.exception import ServerError
1516
from agentrun.utils.log import logger
1617

@@ -453,13 +454,79 @@ async def check_health_async(self):
453454
# Browser API Methods
454455
# ========================================
455456

456-
def get_cdp_url(self, record: Optional[bool] = False):
457-
"""Get CDP WebSocket URL for browser automation."""
458-
return self.data_api.get_cdp_url(record=record)
457+
@overload
458+
def get_cdp_url(
459+
self,
460+
record: Optional[bool] = False,
461+
*,
462+
with_headers: Literal[True],
463+
config: Optional[Config] = None,
464+
) -> Tuple[str, Dict[str, str]]:
465+
...
466+
467+
@overload
468+
def get_cdp_url(
469+
self,
470+
record: Optional[bool] = False,
471+
*,
472+
with_headers: Literal[False] = False,
473+
config: Optional[Config] = None,
474+
) -> str:
475+
...
476+
477+
def get_cdp_url(
478+
self,
479+
record: Optional[bool] = False,
480+
*,
481+
with_headers: bool = False,
482+
config: Optional[Config] = None,
483+
) -> Union[str, Tuple[str, Dict[str, str]]]:
484+
"""Get CDP WebSocket URL for browser automation.
485+
486+
Args:
487+
record: Whether to enable recording / 是否启用录制
488+
with_headers: If True, return (url, headers) tuple with authentication headers.
489+
当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。
490+
config: Optional config override / 可选的配置覆盖
491+
"""
492+
return self.data_api.get_cdp_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload]
459493

460-
def get_vnc_url(self, record: Optional[bool] = False):
461-
"""Get VNC WebSocket URL for live view."""
462-
return self.data_api.get_vnc_url(record=record)
494+
@overload
495+
def get_vnc_url(
496+
self,
497+
record: Optional[bool] = False,
498+
*,
499+
with_headers: Literal[True],
500+
config: Optional[Config] = None,
501+
) -> Tuple[str, Dict[str, str]]:
502+
...
503+
504+
@overload
505+
def get_vnc_url(
506+
self,
507+
record: Optional[bool] = False,
508+
*,
509+
with_headers: Literal[False] = False,
510+
config: Optional[Config] = None,
511+
) -> str:
512+
...
513+
514+
def get_vnc_url(
515+
self,
516+
record: Optional[bool] = False,
517+
*,
518+
with_headers: bool = False,
519+
config: Optional[Config] = None,
520+
) -> Union[str, Tuple[str, Dict[str, str]]]:
521+
"""Get VNC WebSocket URL for live view.
522+
523+
Args:
524+
record: Whether to enable recording / 是否启用录制
525+
with_headers: If True, return (url, headers) tuple with authentication headers.
526+
当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。
527+
config: Optional config override / 可选的配置覆盖
528+
"""
529+
return self.data_api.get_vnc_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload]
463530

464531
def sync_playwright(self, record: Optional[bool] = False):
465532
"""Get synchronous Playwright browser instance."""

‎agentrun/sandbox/__browser_sandbox_async_template.py‎

Lines changed: 74 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,11 @@
66

77
import asyncio
88
import time # noqa: F401
9-
from typing import Optional
9+
from typing import Dict, Literal, Optional, overload, Tuple, Union
1010

1111
from agentrun.sandbox.api import BrowserDataAPI
1212
from agentrun.sandbox.model import TemplateType
13+
from agentrun.utils.config import Config
1314
from agentrun.utils.log import logger
1415

1516
from .sandbox import Sandbox
@@ -78,11 +79,79 @@ def data_api(self):
7879
async def check_health_async(self):
7980
return await self.data_api.check_health_async()
8081

81-
def get_cdp_url(self, record: Optional[bool] = False):
82-
return self.data_api.get_cdp_url(record=record)
82+
@overload
83+
def get_cdp_url(
84+
self,
85+
record: Optional[bool] = False,
86+
*,
87+
with_headers: Literal[True],
88+
config: Optional[Config] = None,
89+
) -> Tuple[str, Dict[str, str]]:
90+
...
91+
92+
@overload
93+
def get_cdp_url(
94+
self,
95+
record: Optional[bool] = False,
96+
*,
97+
with_headers: Literal[False] = False,
98+
config: Optional[Config] = None,
99+
) -> str:
100+
...
101+
102+
def get_cdp_url(
103+
self,
104+
record: Optional[bool] = False,
105+
*,
106+
with_headers: bool = False,
107+
config: Optional[Config] = None,
108+
) -> Union[str, Tuple[str, Dict[str, str]]]:
109+
"""Get CDP WebSocket URL for browser automation.
83110
84-
def get_vnc_url(self, record: Optional[bool] = False):
85-
return self.data_api.get_vnc_url(record=record)
111+
Args:
112+
record: Whether to enable recording / 是否启用录制
113+
with_headers: If True, return (url, headers) tuple with authentication headers.
114+
当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。
115+
config: Optional config override / 可选的配置覆盖
116+
"""
117+
return self.data_api.get_cdp_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload]
118+
119+
@overload
120+
def get_vnc_url(
121+
self,
122+
record: Optional[bool] = False,
123+
*,
124+
with_headers: Literal[True],
125+
config: Optional[Config] = None,
126+
) -> Tuple[str, Dict[str, str]]:
127+
...
128+
129+
@overload
130+
def get_vnc_url(
131+
self,
132+
record: Optional[bool] = False,
133+
*,
134+
with_headers: Literal[False] = False,
135+
config: Optional[Config] = None,
136+
) -> str:
137+
...
138+
139+
def get_vnc_url(
140+
self,
141+
record: Optional[bool] = False,
142+
*,
143+
with_headers: bool = False,
144+
config: Optional[Config] = None,
145+
) -> Union[str, Tuple[str, Dict[str, str]]]:
146+
"""Get VNC WebSocket URL for live view.
147+
148+
Args:
149+
record: Whether to enable recording / 是否启用录制
150+
with_headers: If True, return (url, headers) tuple with authentication headers.
151+
当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。
152+
config: Optional config override / 可选的配置覆盖
153+
"""
154+
return self.data_api.get_vnc_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload]
86155

87156
def sync_playwright(self, record: Optional[bool] = False):
88157
return self.data_api.sync_playwright(record=record)

‎agentrun/sandbox/aio_sandbox.py‎

Lines changed: 74 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,11 @@
1717

1818
import asyncio
1919
import time # noqa: F401
20-
from typing import Optional
20+
from typing import Dict, Literal, Optional, overload, Tuple, Union
2121

2222
from agentrun.sandbox.api.aio_data import AioDataAPI
2323
from agentrun.sandbox.model import CodeLanguage, TemplateType
24+
from agentrun.utils.config import Config
2425
from agentrun.utils.exception import ServerError
2526
from agentrun.utils.log import logger
2627

@@ -817,13 +818,79 @@ def check_health(self):
817818
# Browser API Methods
818819
# ========================================
819820

820-
def get_cdp_url(self, record: Optional[bool] = False):
821-
"""Get CDP WebSocket URL for browser automation."""
822-
return self.data_api.get_cdp_url(record=record)
821+
@overload
822+
def get_cdp_url(
823+
self,
824+
record: Optional[bool] = False,
825+
*,
826+
with_headers: Literal[True],
827+
config: Optional[Config] = None,
828+
) -> Tuple[str, Dict[str, str]]:
829+
...
830+
831+
@overload
832+
def get_cdp_url(
833+
self,
834+
record: Optional[bool] = False,
835+
*,
836+
with_headers: Literal[False] = False,
837+
config: Optional[Config] = None,
838+
) -> str:
839+
...
840+
841+
def get_cdp_url(
842+
self,
843+
record: Optional[bool] = False,
844+
*,
845+
with_headers: bool = False,
846+
config: Optional[Config] = None,
847+
) -> Union[str, Tuple[str, Dict[str, str]]]:
848+
"""Get CDP WebSocket URL for browser automation.
849+
850+
Args:
851+
record: Whether to enable recording / 是否启用录制
852+
with_headers: If True, return (url, headers) tuple with authentication headers.
853+
当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。
854+
config: Optional config override / 可选的配置覆盖
855+
"""
856+
return self.data_api.get_cdp_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload]
823857

824-
def get_vnc_url(self, record: Optional[bool] = False):
825-
"""Get VNC WebSocket URL for live view."""
826-
return self.data_api.get_vnc_url(record=record)
858+
@overload
859+
def get_vnc_url(
860+
self,
861+
record: Optional[bool] = False,
862+
*,
863+
with_headers: Literal[True],
864+
config: Optional[Config] = None,
865+
) -> Tuple[str, Dict[str, str]]:
866+
...
867+
868+
@overload
869+
def get_vnc_url(
870+
self,
871+
record: Optional[bool] = False,
872+
*,
873+
with_headers: Literal[False] = False,
874+
config: Optional[Config] = None,
875+
) -> str:
876+
...
877+
878+
def get_vnc_url(
879+
self,
880+
record: Optional[bool] = False,
881+
*,
882+
with_headers: bool = False,
883+
config: Optional[Config] = None,
884+
) -> Union[str, Tuple[str, Dict[str, str]]]:
885+
"""Get VNC WebSocket URL for live view.
886+
887+
Args:
888+
record: Whether to enable recording / 是否启用录制
889+
with_headers: If True, return (url, headers) tuple with authentication headers.
890+
当为 True 时，返回 (url, headers) 元组，包含鉴权头信息。
891+
config: Optional config override / 可选的配置覆盖
892+
"""
893+
return self.data_api.get_vnc_url(record=record, with_headers=with_headers, config=config) # type: ignore[call-overload]
827894

828895
def sync_playwright(self, record: Optional[bool] = False):
829896
"""Get synchronous Playwright browser instance."""

0 commit comments

Comments
 (0)