diff --git a/src/bedrock_agentcore/tools/browser_client.py b/src/bedrock_agentcore/tools/browser_client.py index 0a5ae0b..fe2897d 100644 --- a/src/bedrock_agentcore/tools/browser_client.py +++ b/src/bedrock_agentcore/tools/browser_client.py @@ -25,6 +25,7 @@ DEFAULT_IDENTIFIER = "aws.browser.v1" DEFAULT_SESSION_TIMEOUT = 3600 DEFAULT_LIVE_VIEW_PRESIGNED_URL_TIMEOUT = 300 +MAX_LIVE_VIEW_PRESIGNED_URL_TIMEOUT = 300 class BrowserClient: @@ -214,15 +215,22 @@ def generate_live_view_url(self, expires: int = DEFAULT_LIVE_VIEW_PRESIGNED_URL_ Args: expires (int, optional): The number of seconds until the pre-signed URL expires. Defaults to DEFAULT_LIVE_VIEW_PRESIGNED_URL_TIMEOUT (300 seconds). + Maximum allowed value is MAX_LIVE_VIEW_PRESIGNED_URL_TIMEOUT seconds. Returns: str: The pre-signed URL for viewing the browser session. Raises: + ValueError: If expires exceeds MAX_LIVE_VIEW_PRESIGNED_URL_TIMEOUT seconds. RuntimeError: If the URL generation fails. """ self.logger.info("Generating live view url...") + if expires > MAX_LIVE_VIEW_PRESIGNED_URL_TIMEOUT: + raise ValueError( + f"Expiry timeout cannot exceed {MAX_LIVE_VIEW_PRESIGNED_URL_TIMEOUT} seconds, got {expires}" + ) + if not self.identifier or not self.session_id: self.start() diff --git a/tests/bedrock_agentcore/tools/test_browser_client.py b/tests/bedrock_agentcore/tools/test_browser_client.py index ff11306..fc491ee 100644 --- a/tests/bedrock_agentcore/tools/test_browser_client.py +++ b/tests/bedrock_agentcore/tools/test_browser_client.py @@ -1,7 +1,11 @@ import datetime from unittest.mock import MagicMock, patch -from bedrock_agentcore.tools.browser_client import BrowserClient, browser_session +from bedrock_agentcore.tools.browser_client import ( + MAX_LIVE_VIEW_PRESIGNED_URL_TIMEOUT, + BrowserClient, + browser_session, +) class TestBrowserClient: @@ -249,7 +253,7 @@ def test_generate_live_view_url(self, mock_get_endpoint, mock_boto3): mock_signer.add_auth.return_value = None # Act - result_url = client.generate_live_view_url(expires=600) + result_url = client.generate_live_view_url(expires=MAX_LIVE_VIEW_PRESIGNED_URL_TIMEOUT) # Assert assert ( @@ -260,9 +264,64 @@ def test_generate_live_view_url(self, mock_get_endpoint, mock_boto3): credentials=mock_frozen_creds, service_name="bedrock-agentcore", region_name="us-west-2", - expires=600, + expires=MAX_LIVE_VIEW_PRESIGNED_URL_TIMEOUT, ) + @patch("bedrock_agentcore.tools.browser_client.boto3") + def test_generate_live_view_url_expires_validation_valid(self, mock_boto3): + # Arrange + client = BrowserClient("us-west-2") + client.identifier = "test-browser-id" + client.session_id = "test-session-id" + + # Mock the dependencies for URL generation + with ( + patch("bedrock_agentcore.tools.browser_client.get_data_plane_endpoint") as mock_get_endpoint, + patch("bedrock_agentcore.tools.browser_client.SigV4QueryAuth") as mock_sigv4_query, + patch("bedrock_agentcore.tools.browser_client.AWSRequest") as mock_aws_request, + ): + mock_get_endpoint.return_value = "https://api.example.com" + + # Mock boto3 session and credentials + mock_boto_session = MagicMock() + mock_credentials = MagicMock() + mock_frozen_creds = MagicMock() + mock_credentials.get_frozen_credentials.return_value = mock_frozen_creds + mock_boto_session.get_credentials.return_value = mock_credentials + mock_boto3.Session.return_value = mock_boto_session + + # Mock the signer and request + mock_signer = MagicMock() + mock_sigv4_query.return_value = mock_signer + + mock_request = MagicMock() + mock_request.url = "https://api.example.com/signed-url" + mock_aws_request.return_value = mock_request + + # Act - test valid expires values + for valid_expires in [1, 150, MAX_LIVE_VIEW_PRESIGNED_URL_TIMEOUT]: + result = client.generate_live_view_url(expires=valid_expires) + # Assert + assert result == "https://api.example.com/signed-url" + + @patch("bedrock_agentcore.tools.browser_client.boto3") + def test_generate_live_view_url_expires_validation_invalid(self, mock_boto3): + # Arrange + client = BrowserClient("us-west-2") + client.identifier = "test-browser-id" + client.session_id = "test-session-id" + + # Act & Assert - test invalid expires values + for invalid_expires in [MAX_LIVE_VIEW_PRESIGNED_URL_TIMEOUT + 1, 500, 1000]: + try: + client.generate_live_view_url(expires=invalid_expires) + raise AssertionError(f"Expected ValueError for expires={invalid_expires}") + except ValueError as e: + expected_msg = ( + f"Expiry timeout cannot exceed {MAX_LIVE_VIEW_PRESIGNED_URL_TIMEOUT} seconds, got {invalid_expires}" + ) + assert expected_msg in str(e) + @patch("bedrock_agentcore.tools.browser_client.boto3") def test_take_control(self, mock_boto3): # Arrange