fix(aws-documentation-mcp-server): Fix get_available_services endpoint (#1796)

* fix(aws-documentation-mcp-server): Fix get_available_services endpoint

A change was recently made to the AWS CN page for available services
and completely changed its format. This change is to update the tool
endpoint function to continue surfacing the available services in
AWS CN.

* fix: Updating tests to cover new functionality introduced in change

* Updating object/dict access patterns to use safer approach

* Using existing approach for returning/logging error messages

---------

Co-authored-by: Jonathan Lim <[email protected]>

Author: JonLim

src/aws-documentation-mcp-server/README.md

@@ -189,3 +189,12 @@ Gets a list of available AWS services in China regions.
 ```python
 get_available_services() -> str
 ```
+
+## Development
+
+For getting started with development on the AWS Documentation MCP server, please refer to the awslabs/mcp DEVELOPER_GUIDE first. Everything below this is specific to AWS Documentation MCP Server development.
+
+### Running tests
+
+Unit tests: `uv run --frozen pytest --cov --cov-branch --cov-report=term-missing`
+Unit tests with integration tests: `uv run --frozen pytest --cov --cov-branch --cov-report=term-missing --run-live`

src/aws-documentation-mcp-server/awslabs/aws_documentation_mcp_server/server_aws_cn.py

@@ -160,6 +160,9 @@ async def get_available_services(
     """
     url_str = 'https://docs.amazonaws.cn/en_us/aws/latest/userguide/services.html'
     url_with_session = f'{url_str}?session={SESSION_UUID}'
+
+    toc_url_str = 'https://docs.amazonaws.cn/en_us/aws/latest/userguide/toc-contents.json'
+    toc_url_with_session = f'{toc_url_str}?session={SESSION_UUID}'
     async with httpx.AsyncClient() as client:
         try:
             response = await client.get(
@@ -168,21 +171,60 @@ async def get_available_services(
                 headers={'User-Agent': DEFAULT_USER_AGENT},
                 timeout=30,
             )
+            # Fetch the Table of Contents in the Services page, which contains the list of supported services
+            toc_response = await client.get(
+                toc_url_with_session,
+                follow_redirects=True,
+                headers={'User-Agent': DEFAULT_USER_AGENT, 'Content-Type': 'application/json'},
+                timeout=30,
+            )
         except httpx.HTTPError as e:
-            error_msg = f'Failed to fetch {url_str}: {str(e)}'
+            error_msg = f'Failed to fetch AWS-CN services page: {str(e)}'
             logger.error(error_msg)
             await ctx.error(error_msg)
             return error_msg
 
         if response.status_code >= 400:
-            error_msg = f'Failed to fetch {url_str} - status code {response.status_code}'
+            error_msg = (
+                f'Failed to fetch AWS-CN services page - status code {response.status_code}'
+            )
             logger.error(error_msg)
             await ctx.error(error_msg)
             return error_msg
 
         page_raw = response.text
         content_type = response.headers.get('content-type', '')
 
+        page_toc_json = toc_response.json()
+        # Expecting a toc JSON object that has a href of 'services.html', which contains all of the AWS Services supported in China
+        # toc_response = { 'contents' : [ { 'title: '', 'href': '', 'contents: [] } ] }
+        services_json = [
+            toc_item.get('contents', [])
+            for toc_item in page_toc_json.get('contents', [])
+            if toc_item.get('href') == 'services.html'
+        ]
+
+        # If toc_response does not have `href: services.html`, and services_json is empty, raise an error so
+        # users can self-solve.
+        if len(services_json) == 0:
+            error_msg = (
+                f'Failed fetching list of available AWS Services, please go to {url_str} directly'
+            )
+            logger.error(error_msg)
+            await ctx.error(error_msg)
+            return error_msg
+
+        # Filtering out 'Services Unsupported in Amazon Web Services in China'
+        formatted_service_titles = ''
+        service_doc_links = [
+            f'[{service.get("title")}](https://docs.amazonaws.cn/en_us/aws/latest/userguide/{service.get("href")})'
+            for service in services_json[0]
+            if 'Services Unsupported' not in service.get('title')
+        ]
+        formatted_service_titles = '\n\n## Services in Amazon Web Services China\n\n' + '\n'.join(
+            [f'- {service_doc_link}' for service_doc_link in service_doc_links]
+        )
+
     if is_html_content(page_raw, content_type):
         content = extract_content_from_html(page_raw)
     else:
@@ -194,7 +236,7 @@ async def get_available_services(
         url_str, content, start_index=0, max_length=MAX_DOCUMENTATION_LENGTH
     )
 
-    return result
+    return result + formatted_service_titles
 
 
 def main():

src/aws-documentation-mcp-server/tests/test_server_aws_cn.py

@@ -117,21 +117,46 @@ async def test_get_available_services(self):
         mock_response.text = '<html><body><h1>AWS Services in China</h1><p>Available services list.</p></body></html>'
         mock_response.headers = {'content-type': 'text/html'}
 
+        mock_toc_response = MagicMock()
+        mock_toc_response.status_code = 200
+        mock_toc_response.json = lambda: {
+            'contents': [
+                {
+                    'title': 'Documentation by Service',
+                    'href': 'services.html',
+                    'contents': [
+                        {'title': 'Amazon Simple Storage Service', 'href': 's3.html'},
+                        {'title': 'Amazon Simple Queue Service', 'href': 'sqs.html'},
+                    ],
+                }
+            ]
+        }
+        mock_toc_response.headers = {'content-type': 'application/json'}
+
         with patch('httpx.AsyncClient.get', new_callable=AsyncMock) as mock_get:
-            mock_get.return_value = mock_response
+            # Set the response for successive calls, first to service.html, second to toc.json
+            mock_get.side_effect = [
+                mock_response,
+                mock_toc_response,
+            ]
+
             with patch(
                 'awslabs.aws_documentation_mcp_server.server_aws_cn.extract_content_from_html'
             ) as mock_extract:
                 mock_extract.return_value = '# AWS Services in China\n\nAvailable services list.'
-
                 result = await get_available_services(ctx)
 
                 assert 'AWS Documentation from' in result
                 assert (
                     'https://docs.amazonaws.cn/en_us/aws/latest/userguide/services.html' in result
                 )
                 assert '# AWS Services in China\n\nAvailable services list.' in result
-                mock_get.assert_called_once()
+                assert 'Amazon Simple Storage Service' in result
+                assert 's3.html' in result
+                assert 'Amazon Simple Queue Service' in result
+                assert 'sqs.html' in result
+
+                assert mock_get.call_count == 2
                 mock_extract.assert_called_once()
                 called_url = mock_get.call_args[0][0]
                 assert '?session=' in called_url
@@ -165,7 +190,7 @@ async def test_get_available_services_status_error(self):
 
             assert 'Failed to fetch' in result
             assert 'status code 404' in result
-            mock_get.assert_called_once()
+            assert mock_get.call_count == 2
 
     @pytest.mark.asyncio
     async def test_get_available_services_non_html(self):
@@ -177,8 +202,29 @@ async def test_get_available_services_non_html(self):
         mock_response.text = 'Plain text content'
         mock_response.headers = {'content-type': 'text/plain'}
 
+        mock_toc_response = MagicMock()
+        mock_toc_response.status_code = 200
+        mock_toc_response.json = lambda: {
+            'contents': [
+                {
+                    'title': 'Documentation by Service',
+                    'href': 'services.html',
+                    'contents': [
+                        {'title': 'Amazon Simple Storage Service', 'href': 's3.html'},
+                        {'title': 'Amazon Simple Queue Service', 'href': 'sqs.html'},
+                    ],
+                }
+            ]
+        }
+        mock_toc_response.headers = {'content-type': 'application/json'}
+
         with patch('httpx.AsyncClient.get', new_callable=AsyncMock) as mock_get:
-            mock_get.return_value = mock_response
+            # Set the response for successive calls, first to service.html, second to toc.json
+            mock_get.side_effect = [
+                mock_response,
+                mock_toc_response,
+            ]
+
             with patch(
                 'awslabs.aws_documentation_mcp_server.server_aws_cn.is_html_content'
             ) as mock_is_html:
@@ -188,9 +234,51 @@ async def test_get_available_services_non_html(self):
 
                 assert 'AWS Documentation from' in result
                 assert 'Plain text content' in result
-                mock_get.assert_called_once()
+                assert mock_get.call_count == 2
                 mock_is_html.assert_called_once()
 
+    @pytest.mark.asyncio
+    async def test_get_available_services_key_error(self):
+        """Test getting available services in AWS China."""
+        ctx = MockContext()
+
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.text = '<html><body><h1>AWS Services in China</h1><p>Available services list.</p></body></html>'
+        mock_response.headers = {'content-type': 'text/html'}
+
+        mock_toc_response = MagicMock()
+        mock_toc_response.status_code = 200
+        mock_toc_response.json = lambda: {
+            'contents': [
+                {
+                    'title': 'Welcome',
+                    'href': 'introduction.html',
+                }
+            ]
+        }
+        mock_toc_response.headers = {'content-type': 'application/json'}
+
+        with patch('httpx.AsyncClient.get', new_callable=AsyncMock) as mock_get:
+            # Set the response for successive calls, first to service.html, second to toc.json
+            mock_get.side_effect = [
+                mock_response,
+                mock_toc_response,
+            ]
+
+            with patch(
+                'awslabs.aws_documentation_mcp_server.server_aws_cn.extract_content_from_html'
+            ) as mock_extract:
+                mock_extract.return_value = '# AWS Services in China\n\nAvailable services list.'
+                result = await get_available_services(ctx)
+
+                assert 'Failed fetching list of available AWS Services, please go to' in result
+
+                assert mock_get.call_count == 2
+                mock_extract.assert_not_called()
+                called_url = mock_get.call_args[0][0]
+                assert '?session=' in called_url
+
 
 class TestMain:
     """Tests for the main function."""