⚡️ Speed up method `Server.get_capabilities` by 9% by codeflash-ai[bot] · Pull Request #7 · Saga4/python-sdk

codeflash-ai · Apr 29, 2025

📄 9% (0.09x) speedup for `Server.get_capabilities` in `src/mcp/server/lowlevel/server.py`

⏱️ Runtime : 284 microseconds → 262 microseconds (best of 199 runs)

📝 Explanation and details

Here is a highly optimized rewrite of your code, focusing on get_capabilities (the bottleneck, >76% of time is spent in return value construction), with memory and runtime optimizations. These include.

Use __slots__ for server class to save memory.
Precompute and cache all field values and handler presence checks in local variables.
Only construct capability objects and pass them to ServerCapabilities if a handler actually exists (reduces unnecessary assignments).
Avoid creating unnecessary intermediate variables.
Compute all in self.request_handlers lookups once.
Reduce types.ServerCapabilities construction to the minimal needed.
Move logger.debug(...) outside __init__ to after all attributes are set.
Minimal allocation in the fast path.

No imports or API signatures were changed. All comments preserved and unchanged.

This version builds the exact same objects, preserves all functional behavior, and accelerates the most time/spending path, especially by reducing unnecessary work, minimizing dictionary lookups, and using argument unpacking to only create and pass populated fields.

Let me know if you want further performance tuning or Cythonization!

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 30 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests Details

import logging
from collections.abc import Awaitable, Callable
from contextlib import AbstractAsyncContextManager
from typing import Any, Generic, TypeVar
from unittest.mock import MagicMock

import mcp.types as types
# imports
import pytest  # used for our unit tests
from src.mcp.server.lowlevel.server import Server

logger = logging.getLogger(__name__)

LifespanResultT = TypeVar("LifespanResultT")
from src.mcp.server.lowlevel.server import Server


async def _ping_handler(request: types.PingRequest) -> types.ServerResult:
    return types.ServerResult(types.EmptyResult())


# Mock classes for testing
class NotificationOptions:
    def __init__(self, prompts_changed=False, resources_changed=False, tools_changed=False):
        self.prompts_changed = prompts_changed
        self.resources_changed = resources_changed
        self.tools_changed = tools_changed


# unit tests
def test_no_handlers_registered():
    """Test when no handlers are registered."""
    server = Server(name="TestServer")
    server.request_handlers = {}  # Clear all handlers
    codeflash_output = server.get_capabilities(NotificationOptions(), {}); capabilities = codeflash_output

def test_single_handler_registered():
    """Test with a single handler registered."""
    server = Server(name="TestServer")
    server.request_handlers = {types.ListPromptsRequest: MagicMock()}
    codeflash_output = server.get_capabilities(NotificationOptions(prompts_changed=True), {}); capabilities = codeflash_output

def test_multiple_handlers_registered():
    """Test with multiple handlers registered."""
    server = Server(name="TestServer")
    server.request_handlers = {
        types.ListPromptsRequest: MagicMock(),
        types.ListResourcesRequest: MagicMock(),
        types.ListToolsRequest: MagicMock(),
        types.SetLevelRequest: MagicMock()
    }
    codeflash_output = server.get_capabilities(NotificationOptions(prompts_changed=True, resources_changed=True, tools_changed=True), {}); capabilities = codeflash_output

def test_unknown_request_type():
    """Test with an unknown request type."""
    server = Server(name="TestServer")
    server.request_handlers = {MagicMock(): MagicMock()}  # Unknown type
    codeflash_output = server.get_capabilities(NotificationOptions(), {}); capabilities = codeflash_output

def test_empty_experimental_capabilities():
    """Test with empty experimental capabilities."""
    server = Server(name="TestServer")
    codeflash_output = server.get_capabilities(NotificationOptions(), {}); capabilities = codeflash_output

def test_complex_experimental_capabilities():
    """Test with complex experimental capabilities."""
    server = Server(name="TestServer")
    experimental_capabilities = {"featureX": {"enabled": True, "version": 1.0}}
    codeflash_output = server.get_capabilities(NotificationOptions(), experimental_capabilities); capabilities = codeflash_output

def test_notification_options_all_false():
    """Test with all notification options set to False."""
    server = Server(name="TestServer")
    server.request_handlers = {types.ListPromptsRequest: MagicMock()}
    codeflash_output = server.get_capabilities(NotificationOptions(), {}); capabilities = codeflash_output

def test_notification_options_mixed_flags():
    """Test with mixed notification options."""
    server = Server(name="TestServer")
    server.request_handlers = {types.ListPromptsRequest: MagicMock(), types.ListResourcesRequest: MagicMock()}
    codeflash_output = server.get_capabilities(NotificationOptions(prompts_changed=True), {}); capabilities = codeflash_output

def test_large_number_of_handlers():
    """Test with a large number of handlers."""
    server = Server(name="TestServer")
    server.request_handlers = {MagicMock(): MagicMock() for _ in range(100)}  # Simulate many handlers
    codeflash_output = server.get_capabilities(NotificationOptions(), {}); capabilities = codeflash_output

def test_complex_notification_options():
    """Test with complex notification options."""
    server = Server(name="TestServer")
    server.request_handlers = {types.ListPromptsRequest: MagicMock()}
    complex_options = NotificationOptions(prompts_changed=True, resources_changed=True, tools_changed=True)
    codeflash_output = server.get_capabilities(complex_options, {}); capabilities = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import logging
from collections.abc import Awaitable, Callable
from contextlib import AbstractAsyncContextManager
from typing import Any, Generic, TypeVar
from unittest.mock import MagicMock

import mcp.types as types
# imports
import pytest  # used for our unit tests
from src.mcp.server.lowlevel.server import Server

logger = logging.getLogger(__name__)

LifespanResultT = TypeVar("LifespanResultT")
from src.mcp.server.lowlevel.server import Server


async def _ping_handler(request: types.PingRequest) -> types.ServerResult:
    return types.ServerResult(types.EmptyResult())


# Mock classes for testing
class NotificationOptions:
    def __init__(self, prompts_changed=False, resources_changed=False, tools_changed=False):
        self.prompts_changed = prompts_changed
        self.resources_changed = resources_changed
        self.tools_changed = tools_changed

class MockTypes:
    class ServerCapabilities:
        def __init__(self, prompts, resources, tools, logging, experimental):
            self.prompts = prompts
            self.resources = resources
            self.tools = tools
            self.logging = logging
            self.experimental = experimental

    class PromptsCapability:
        def __init__(self, listChanged):
            self.listChanged = listChanged

    class ResourcesCapability:
        def __init__(self, subscribe, listChanged):
            self.subscribe = subscribe
            self.listChanged = listChanged

    class ToolsCapability:
        def __init__(self, listChanged):
            self.listChanged = listChanged

    class LoggingCapability:
        pass

    class PingRequest:
        pass

    class ListPromptsRequest:
        pass

    class ListResourcesRequest:
        pass

    class ListToolsRequest:
        pass

    class SetLevelRequest:
        pass

    class ServerResult:
        def __init__(self, result):
            self.result = result

    class EmptyResult:
        pass

types = MockTypes()

# unit tests
def test_no_handlers():
    """Test with no handlers registered."""
    server = Server("TestServer")
    notification_options = NotificationOptions()
    experimental_capabilities = {}
    codeflash_output = server.get_capabilities(notification_options, experimental_capabilities); capabilities = codeflash_output

def test_single_prompt_handler():
    """Test with only ListPromptsRequest handler registered."""
    server = Server("TestServer")
    server.request_handlers[types.ListPromptsRequest] = MagicMock()
    notification_options = NotificationOptions(prompts_changed=True)
    experimental_capabilities = {}
    codeflash_output = server.get_capabilities(notification_options, experimental_capabilities); capabilities = codeflash_output

def test_all_handlers():
    """Test with all handlers registered."""
    server = Server("TestServer")
    server.request_handlers.update({
        types.ListPromptsRequest: MagicMock(),
        types.ListResourcesRequest: MagicMock(),
        types.ListToolsRequest: MagicMock(),
        types.SetLevelRequest: MagicMock(),
    })
    notification_options = NotificationOptions(prompts_changed=True, resources_changed=True, tools_changed=True)
    experimental_capabilities = {"featureX": {"enabled": True}}
    codeflash_output = server.get_capabilities(notification_options, experimental_capabilities); capabilities = codeflash_output

def test_large_experimental_capabilities():
    """Test with a large number of experimental capabilities."""
    server = Server("TestServer")
    experimental_capabilities = {f"feature{i}": {"enabled": True} for i in range(1000)}
    codeflash_output = server.get_capabilities(NotificationOptions(), experimental_capabilities); capabilities = codeflash_output

def test_invalid_handler_type():
    """Test with an invalid handler type."""
    server = Server("TestServer")
    server.request_handlers["InvalidRequest"] = MagicMock()
    codeflash_output = server.get_capabilities(NotificationOptions(), {}); capabilities = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from src.mcp.server.lowlevel.server import Server

To edit these changes git checkout codeflash/optimize-Server.get_capabilities-ma2z0cak and push.

Here is a highly optimized rewrite of your code, focusing on `get_capabilities` (the bottleneck, >76% of time is spent in return value construction), with memory and runtime optimizations. These include. - Use `__slots__` for server class to save memory. - Precompute and cache all field values and handler presence checks in local variables. - Only construct capability objects and pass them to `ServerCapabilities` if a handler actually exists (reduces unnecessary assignments). - Avoid creating unnecessary intermediate variables. - Compute all `in self.request_handlers` lookups once. - Reduce `types.ServerCapabilities` construction to the minimal needed. - Move `logger.debug(...)` outside `__init__` to after all attributes are set. - Minimal allocation in the fast path. **No imports or API signatures were changed. All comments preserved and unchanged.** This version builds the exact same objects, preserves all functional behavior, and accelerates the most time/spending path, especially by reducing unnecessary work, minimizing dictionary lookups, and using argument unpacking to only create and pass populated fields. Let me know if you want further performance tuning or Cythonization!

codeflash-ai bot added the ⚡️ codeflash Optimization PR opened by Codeflash AI label Apr 29, 2025

codeflash-ai bot requested a review from Saga4 April 29, 2025 20:39

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Comments

⚡️ Speed up method `Server.get_capabilities` by 9%#7

⚡️ Speed up method `Server.get_capabilities` by 9%#7
codeflash-ai[bot] wants to merge 1 commit intomainSaga4/python-sdk:mainfrom
codeflash/optimize-Server.get_capabilities-ma2z0cakSaga4/python-sdk:codeflash/optimize-Server.get_capabilities-ma2z0cakCopy head branch name to clipboard

codeflash-ai bot commented Apr 29, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

0 participants

Search code, repositories, users, issues, pull requests...

Comments

Conversation

codeflash-ai bot commented Apr 29, 2025

📄 9% (0.09x) speedup for Server.get_capabilities in src/mcp/server/lowlevel/server.py

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

0 participants

📄 9% (0.09x) speedup for `Server.get_capabilities` in `src/mcp/server/lowlevel/server.py`