Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 23 additions & 18 deletions src/bedrock_agentcore/tools/browser_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@ def start(
identifier: Optional[str] = DEFAULT_IDENTIFIER,
name: Optional[str] = None,
session_timeout_seconds: Optional[int] = DEFAULT_SESSION_TIMEOUT,
viewport: Optional[Dict[str, int]] = None,
) -> str:
"""Start a browser sandbox session.

Expand All @@ -111,19 +112,31 @@ def start(
will be generated using a UUID.
session_timeout_seconds (Optional[int]): The timeout for the session in seconds.
Defaults to DEFAULT_TIMEOUT.
description (Optional[str]): A description for this session.
Defaults to an empty string.
viewport (Optional[Dict[str, int]]): The viewport dimensions for the browser.
Should be a dict with 'width' and 'height' keys (e.g., {'width': 1920, 'height': 1080}).
If not provided, the service default viewport will be used.

Returns:
str: The session ID of the newly created session.

Example:
>>> client = BrowserClient('us-west-2')
>>> session_id = client.start(viewport={'width': 1920, 'height': 1080})
"""
self.logger.info("Starting browser session...")

response = self.client.start_browser_session(
browserIdentifier=identifier,
name=name or f"browser-session-{uuid.uuid4().hex[:8]}",
sessionTimeoutSeconds=session_timeout_seconds,
)
# Build the request parameters
request_params = {
"browserIdentifier": identifier,
"name": name or f"browser-session-{uuid.uuid4().hex[:8]}",
"sessionTimeoutSeconds": session_timeout_seconds,
}

# Add viewport if provided
if viewport is not None:
request_params["viewPort"] = viewport

response = self.client.start_browser_session(**request_params)

self.identifier = response["browserIdentifier"]
self.session_id = response["sessionId"]
Expand Down Expand Up @@ -306,26 +319,18 @@ def _update_browser_stream(self, identifier: str, session_id: str, stream_status


@contextmanager
def browser_session(region: str) -> Generator[BrowserClient, None, None]:
def browser_session(region: str, viewport: Optional[Dict[str, int]] = None) -> Generator[BrowserClient, None, None]:
"""Context manager for creating and managing a browser sandbox session.

This context manager handles creating a client, starting a session, and
ensuring the session is properly cleaned up when the context exits.

Args:
region (str): The AWS region to use for the Browser service.
viewport (Optional[Dict[str, int]]): The viewport dimensions for the browser.

Yields:
BrowserClient: An initialized and started browser client.

Example:
>>> with browser_session('us-west-2') as client:
... browser = client.get_browser_obj()
... page = browser.new_page()
... page.goto('https://example.com')
"""
client = BrowserClient(region)
client.start()
client.start(viewport=viewport)

try:
yield client
Expand Down
47 changes: 46 additions & 1 deletion tests/bedrock_agentcore/tools/test_browser_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,10 +71,12 @@ def test_start_with_defaults(self, mock_uuid4, mock_boto3):
assert client.session_id == "session-123"

@patch("bedrock_agentcore.tools.browser_client.boto3")
def test_start_with_custom_params(self, mock_boto3):
@patch("bedrock_agentcore.tools.browser_client.uuid.uuid4")
def test_start_with_custom_params(self, mock_uuid4, mock_boto3):
# Arrange
mock_client = MagicMock()
mock_boto3.client.return_value = mock_client
mock_uuid4.return_value.hex = "12345678abcdef"

client = BrowserClient("us-west-2")
mock_response = {"browserIdentifier": "custom.browser", "sessionId": "custom-session-123"}
Expand Down Expand Up @@ -361,3 +363,46 @@ def test_release_control(self, mock_boto3):
sessionId="test-session-id",
streamUpdate={"automationStreamUpdate": {"streamStatus": "ENABLED"}},
)

@patch("bedrock_agentcore.tools.browser_client.boto3")
@patch("bedrock_agentcore.tools.browser_client.uuid.uuid4")
def test_start_with_viewport(self, mock_uuid4, mock_boto3):
# Arrange
mock_client = MagicMock()
mock_boto3.client.return_value = mock_client
mock_uuid4.return_value.hex = "12345678abcdef"

client = BrowserClient("us-west-2")
mock_response = {"browserIdentifier": "aws.browser.v1", "sessionId": "session-123"}
mock_client.start_browser_session.return_value = mock_response
viewport = {"width": 1920, "height": 1080}

# Act
session_id = client.start(viewport=viewport)

# Assert
mock_client.start_browser_session.assert_called_once_with(
browserIdentifier="aws.browser.v1",
name="browser-session-12345678",
sessionTimeoutSeconds=3600,
viewPort=viewport,
)
assert session_id == "session-123"
assert client.identifier == "aws.browser.v1"
assert client.session_id == "session-123"

@patch("bedrock_agentcore.tools.browser_client.BrowserClient")
def test_browser_session_context_manager_with_viewport(self, mock_client_class):
# Arrange
mock_client = MagicMock()
mock_client_class.return_value = mock_client
viewport = {"width": 1280, "height": 720}

# Act
with browser_session("us-west-2", viewport=viewport):
pass

# Assert
mock_client_class.assert_called_once_with("us-west-2")
mock_client.start.assert_called_once_with(viewport=viewport)
mock_client.stop.assert_called_once()
5 changes: 5 additions & 0 deletions tests_integ/tools/test_browser.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,3 +10,8 @@

client.take_control()
client.release_control()

with browser_session("us-west-2", viewport={"width": 1280, "height": 720}) as client:
assert client.session_id is not None
url, headers = client.generate_ws_headers()
assert url.startswith("wss")
Loading