fix: pass-through server capbilities (#264)

Author: arnewouters

mcp_proxy_for_aws/middleware/initialize_middleware.py

@@ -31,6 +31,23 @@ def __init__(self, client_factory: AWSMCPProxyClientFactory) -> None:
         super().__init__()
         self._client_factory = client_factory
 
+    def _overwrite_init_options(
+        self, context: MiddlewareContext, init_result: mt.InitializeResult
+    ):
+        """Overwrite the session's _init_options with the backend server's info.
+
+        The MCP SDK builds the InitializeResult from session._init_options
+        inside call_next. By modifying _init_options before call_next runs,
+        the response sent to the client will contain the backend server's
+        info instead of the proxy's defaults.
+        """
+        fastmcp_ctx = context.fastmcp_context
+        if fastmcp_ctx is None or fastmcp_ctx._session is None:
+            logger.debug('No session available, skipping init_options overwrite.')
+            return
+
+        fastmcp_ctx._session._init_options.capabilities = init_result.capabilities
+
     @override
     async def on_initialize(
         self,
@@ -59,6 +76,12 @@ async def on_initialize(
                 # the list_tool call will require the client to be connected again, so the mcp error
                 # will be displayed in the q cli logs.
                 await client._connect()
+
+            # Overwrite the proxy's init_options with the backend server's info
+            # so the InitializeResult sent to the client reflects the backend.
+            if client.initialize_result is not None:
+                self._overwrite_init_options(context, client.initialize_result)
+
             return await call_next(context)
         except Exception:
             logger.exception('Initialize failed in middleware.')

tests/unit/test_initialize_middleware.py

@@ -35,15 +35,27 @@ async def test_on_initialize_connects_client():
     """Test that on_initialize calls client._connect()."""
     mock_client = Mock()
     mock_client._connect = AsyncMock()
+    mock_client.initialize_result = mt.InitializeResult(
+        protocolVersion='2024-11-05',
+        capabilities=mt.ServerCapabilities(),
+        serverInfo=mt.Implementation(name='backend-server', version='2.0'),
+    )
 
     mock_factory = Mock()
     mock_factory.set_init_params = Mock()
     mock_factory.get_client = AsyncMock(return_value=mock_client)
 
     middleware = InitializeMiddleware(mock_factory)
 
+    mock_init_options = Mock()
+    mock_session = Mock()
+    mock_session._init_options = mock_init_options
+    mock_fastmcp_ctx = Mock()
+    mock_fastmcp_ctx._session = mock_session
+
     mock_context = Mock()
     mock_context.message = create_initialize_request('test-client')
+    mock_context.fastmcp_context = mock_fastmcp_ctx
 
     mock_call_next = AsyncMock()
 
@@ -54,6 +66,9 @@ async def test_on_initialize_connects_client():
     mock_client._connect.assert_called_once()
     mock_call_next.assert_called_once_with(mock_context)
 
+    # Verify init_options capabilities were overwritten with backend server info
+    assert mock_init_options.capabilities == mt.ServerCapabilities()
+
 
 @pytest.mark.asyncio
 async def test_on_initialize_fails_if_connect_fails():
@@ -94,6 +109,7 @@ async def test_on_initialize_skips_connect_for_special_clients(client_name):
     """Test that on_initialize skips _connect() for Kiro CLI and Q Dev CLI."""
     mock_client = Mock()
     mock_client._connect = AsyncMock()
+    mock_client.initialize_result = None
 
     mock_factory = Mock()
     mock_factory.set_init_params = Mock()
@@ -110,3 +126,116 @@ async def test_on_initialize_skips_connect_for_special_clients(client_name):
 
     mock_client._connect.assert_not_called()
     mock_call_next.assert_called_once_with(mock_context)
+
+
+@pytest.mark.asyncio
+async def test_on_initialize_overwrites_init_options_with_backend_info():
+    """Test that on_initialize overwrites session init_options with backend server info."""
+    backend_capabilities = mt.ServerCapabilities(
+        logging=mt.LoggingCapability(),
+    )
+    backend_result = mt.InitializeResult(
+        protocolVersion='2024-11-05',
+        capabilities=backend_capabilities,
+        serverInfo=mt.Implementation(name='backend-mcp', version='3.1'),
+    )
+
+    mock_client = Mock()
+    mock_client._connect = AsyncMock()
+    mock_client.initialize_result = backend_result
+
+    mock_factory = Mock()
+    mock_factory.set_init_params = Mock()
+    mock_factory.get_client = AsyncMock(return_value=mock_client)
+
+    middleware = InitializeMiddleware(mock_factory)
+
+    mock_init_options = Mock()
+    mock_init_options.server_name = 'proxy-name'
+    mock_init_options.server_version = '1.0'
+    mock_init_options.capabilities = mt.ServerCapabilities()
+    mock_session = Mock()
+    mock_session._init_options = mock_init_options
+    mock_fastmcp_ctx = Mock()
+    mock_fastmcp_ctx._session = mock_session
+
+    mock_context = Mock()
+    mock_context.message = create_initialize_request('test-client')
+    mock_context.fastmcp_context = mock_fastmcp_ctx
+
+    mock_call_next = AsyncMock()
+
+    await middleware.on_initialize(mock_context, mock_call_next)
+
+    assert mock_init_options.capabilities == backend_capabilities
+
+
+@pytest.mark.asyncio
+async def test_on_initialize_disables_prompts_and_resources():
+    """Test that prompts and resources capabilities are disabled even if backend supports them."""
+    backend_capabilities = mt.ServerCapabilities(
+        tools=mt.ToolsCapability(),
+        prompts=mt.PromptsCapability(),
+        resources=mt.ResourcesCapability(),
+    )
+    backend_result = mt.InitializeResult(
+        protocolVersion='2024-11-05',
+        capabilities=backend_capabilities,
+        serverInfo=mt.Implementation(name='backend', version='1.0'),
+    )
+
+    mock_client = Mock()
+    mock_client._connect = AsyncMock()
+    mock_client.initialize_result = backend_result
+
+    mock_factory = Mock()
+    mock_factory.set_init_params = Mock()
+    mock_factory.get_client = AsyncMock(return_value=mock_client)
+
+    middleware = InitializeMiddleware(mock_factory)
+
+    mock_init_options = Mock()
+    mock_session = Mock()
+    mock_session._init_options = mock_init_options
+    mock_fastmcp_ctx = Mock()
+    mock_fastmcp_ctx._session = mock_session
+
+    mock_context = Mock()
+    mock_context.message = create_initialize_request('test-client')
+    mock_context.fastmcp_context = mock_fastmcp_ctx
+
+    mock_call_next = AsyncMock()
+
+    await middleware.on_initialize(mock_context, mock_call_next)
+
+    assert mock_init_options.capabilities.prompts is not None
+    assert mock_init_options.capabilities.resources is not None
+    assert mock_init_options.capabilities.tools is not None
+
+
+@pytest.mark.asyncio
+async def test_on_initialize_skips_overwrite_when_no_session():
+    """Test that overwrite is skipped when no session is available."""
+    mock_client = Mock()
+    mock_client._connect = AsyncMock()
+    mock_client.initialize_result = mt.InitializeResult(
+        protocolVersion='2024-11-05',
+        capabilities=mt.ServerCapabilities(),
+        serverInfo=mt.Implementation(name='backend', version='1.0'),
+    )
+
+    mock_factory = Mock()
+    mock_factory.set_init_params = Mock()
+    mock_factory.get_client = AsyncMock(return_value=mock_client)
+
+    middleware = InitializeMiddleware(mock_factory)
+
+    mock_context = Mock()
+    mock_context.message = create_initialize_request('test-client')
+    mock_context.fastmcp_context = None
+
+    mock_call_next = AsyncMock()
+
+    # Should not raise, just skip overwrite
+    await middleware.on_initialize(mock_context, mock_call_next)
+    mock_call_next.assert_called_once_with(mock_context)