rename result_cache class and allow other implementations for enterprise scenarios where memory cache may not be sufficient

Author: davemssavage

src/mcp/server/lowlevel/result_cache.py …server/lowlevel/async_request_manager.pysrc/mcp/server/lowlevel/result_cache.py renamed to src/mcp/server/lowlevel/async_request_manager.py src/mcp/server/lowlevel/result_cache.py renamed to src/mcp/server/lowlevel/async_request_manager.py

@@ -20,6 +20,36 @@
 logger = getLogger(__name__)
 
 
+class AsyncRequestManager:
+    async def __aenter__(self): ...
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool | None: ...
+    async def start_call(
+        self,
+        call: Callable[[types.CallToolRequest], Awaitable[types.ServerResult]],
+        req: types.CallToolAsyncRequest,
+        ctx: RequestContext[ServerSession, Any, Any],
+    ) -> types.CallToolAsyncResult: ...
+    async def join_call(
+        self,
+        req: types.JoinCallToolAsyncRequest,
+        ctx: RequestContext[ServerSession, Any, Any],
+    ) -> types.CallToolAsyncResult: ...
+    async def cancel(self, notification: types.CancelToolAsyncNotification) -> None: ...
+    async def get_result(
+        self, req: types.GetToolAsyncResultRequest
+    ) -> types.CallToolResult: ...
+
+    async def notification_hook(
+        self, session: ServerSession, notification: types.ServerNotification
+    ) -> None: ...
+    async def session_close_hook(self, session: ServerSession): ...
+
+
 @dataclass
 class InProgress:
     token: str
@@ -40,23 +70,18 @@ def is_expired(self):
             return int(self.timer()) > self.keep_alive_start + self.keep_alive
 
 
-class ResultCache:
+class SimpleInMemoryAsyncRequestManager(AsyncRequestManager):
     """
     Note this class is a work in progress
     Its purpose is to act as a central point for managing in progress
     async calls, allowing multiple clients to join and receive progress
     updates, get results and/or cancel in progress calls
-    TODO CRITICAL not obvious user context will be passed to background thread
-    add tests to assert behaviour with authenticated calls
     TODO MAJOR needs a lot more testing around edge cases/failure scenarios
     TODO MAJOR decide if async.Locks are required for integrity of internal
     data structures
-    TODO ENHANCEMENT externalise cachetools to allow for other implementations
-    e.g. redis etal for production scenarios
     TODO ENHANCEMENT may need to add an authorisation layer to decide if
     a user is allowed to get/join/cancel an existing async call current
     simple logic only allows same user to perform these tasks
-    TODO TRIVIAL name is probably not quite right, more of a result broker?
     """
 
     _in_progress: dict[types.AsyncToken, InProgress]
@@ -178,7 +203,9 @@ async def cancel(self, notification: types.CancelToolAsyncNotification) -> None:
                     f"from {user_context.get()}"
                 )
 
-    async def get_result(self, req: types.GetToolAsyncResultRequest):
+    async def get_result(
+        self, req: types.GetToolAsyncResultRequest
+    ) -> types.CallToolResult:
         logger.debug("Getting result")
         async_token = req.params.token
         in_progress = self._in_progress.get(async_token)

src/mcp/server/lowlevel/server.py

@@ -80,8 +80,11 @@ async def main():
 from typing_extensions import TypeVar
 
 import mcp.types as types
+from mcp.server.lowlevel.async_request_manager import (
+    AsyncRequestManager,
+    SimpleInMemoryAsyncRequestManager,
+)
 from mcp.server.lowlevel.helper_types import ReadResourceContents
-from mcp.server.lowlevel.result_cache import ResultCache
 from mcp.server.models import InitializationOptions
 from mcp.server.session import ServerSession
 from mcp.server.stdio import stdio_server as stdio_server
@@ -136,8 +139,9 @@ def __init__(
             [Server[LifespanResultT, RequestT]],
             AbstractAsyncContextManager[LifespanResultT],
         ] = lifespan,
-        max_cache_size: int = 1000,
-        max_cache_ttl: int = 60,
+        async_request_manager: AsyncRequestManager = SimpleInMemoryAsyncRequestManager(
+            max_size=1000, max_keep_alive=60
+        ),
     ):
         self.name = name
         self.version = version
@@ -148,7 +152,7 @@ def __init__(
         ] = {
             types.PingRequest: _ping_handler,
         }
-        self.result_cache = ResultCache(max_cache_size, max_cache_ttl)
+        self.async_request_manager = async_request_manager
         self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {}
         self.notification_options = NotificationOptions()
         logger.debug("Initializing server %r", name)
@@ -432,19 +436,19 @@ async def handler(req: types.CallToolRequest):
 
             async def async_call_handler(req: types.CallToolAsyncRequest):
                 ctx = request_ctx.get()
-                result = await self.result_cache.start_call(handler, req, ctx)
+                result = await self.async_request_manager.start_call(handler, req, ctx)
                 return types.ServerResult(result)
 
             async def async_join_handler(req: types.JoinCallToolAsyncRequest):
                 ctx = request_ctx.get()
-                result = await self.result_cache.join_call(req, ctx)
+                result = await self.async_request_manager.join_call(req, ctx)
                 return types.ServerResult(result)
 
             async def async_cancel_handler(req: types.CancelToolAsyncNotification):
-                await self.result_cache.cancel(req)
+                await self.async_request_manager.cancel(req)
 
             async def async_result_handler(req: types.GetToolAsyncResultRequest):
-                result = await self.result_cache.get_result(req)
+                result = await self.async_request_manager.get_result(req)
                 return types.ServerResult(result)
 
             self.request_handlers[types.CallToolRequest] = handler
@@ -534,11 +538,11 @@ async def run(
                     write_stream,
                     initialization_options,
                     stateless=stateless,
-                    notification_hook=self.result_cache.notification_hook,
-                    session_close_hook=self.result_cache.session_close_hook,
+                    notification_hook=self.async_request_manager.notification_hook,
+                    session_close_hook=self.async_request_manager.session_close_hook,
                 )
             )
-            await stack.enter_async_context(self.result_cache)
+            await stack.enter_async_context(self.async_request_manager)
 
             async with anyio.create_task_group() as tg:
                 async for message in session.incoming_messages:

…sts/server/lowlevel/test_result_cache.py …r/lowlevel/test_async_request_manager.pytests/server/lowlevel/test_result_cache.py renamed to tests/server/lowlevel/test_async_request_manager.py tests/server/lowlevel/test_result_cache.py renamed to tests/server/lowlevel/test_async_request_manager.py

@@ -8,7 +8,7 @@
 from mcp.server.auth.middleware.auth_context import (
     auth_context_var as user_context,
 )
-from mcp.server.lowlevel.result_cache import ResultCache
+from mcp.server.lowlevel.async_request_manager import SimpleInMemoryAsyncRequestManager
 
 
 @pytest.mark.anyio
@@ -27,7 +27,7 @@ async def test_call(call: types.CallToolRequest) -> types.ServerResult:
     mock_session = AsyncMock()
     mock_context = Mock()
     mock_context.session = mock_session
-    result_cache = ResultCache(max_size=1, max_keep_alive=1)
+    result_cache = SimpleInMemoryAsyncRequestManager(max_size=1, max_keep_alive=1)
     async with AsyncExitStack() as stack:
         await stack.enter_async_context(result_cache)
         async_call_ref = await result_cache.start_call(
@@ -73,7 +73,7 @@ async def test_call(call: types.CallToolRequest) -> types.ServerResult:
 
     mock_context_2.session = mock_session_2
 
-    result_cache = ResultCache(max_size=1, max_keep_alive=1)
+    result_cache = SimpleInMemoryAsyncRequestManager(max_size=1, max_keep_alive=1)
     async with AsyncExitStack() as stack:
         await stack.enter_async_context(result_cache)
         async_call_ref = await result_cache.start_call(
@@ -160,7 +160,7 @@ async def slow_call(call: types.CallToolRequest) -> types.ServerResult:
     mock_context_1 = Mock()
     mock_context_1.session = mock_session_1
 
-    result_cache = ResultCache(max_size=1, max_keep_alive=1)
+    result_cache = SimpleInMemoryAsyncRequestManager(max_size=1, max_keep_alive=1)
     async with AsyncExitStack() as stack:
         await stack.enter_async_context(result_cache)
         async_call_ref = await result_cache.start_call(
@@ -228,7 +228,7 @@ async def test_call(call: types.CallToolRequest) -> types.ServerResult:
 
     mock_context_2.session = mock_session_2
 
-    result_cache = ResultCache(max_size=1, max_keep_alive=10)
+    result_cache = SimpleInMemoryAsyncRequestManager(max_size=1, max_keep_alive=10)
     async with AsyncExitStack() as stack:
         await stack.enter_async_context(result_cache)
         async_call_ref = await result_cache.start_call(
@@ -307,7 +307,9 @@ async def test_call(call: types.CallToolRequest) -> types.ServerResult:
     def test_timer():
         return time
 
-    result_cache = ResultCache(max_size=1, max_keep_alive=1, timer=test_timer)
+    result_cache = SimpleInMemoryAsyncRequestManager(
+        max_size=1, max_keep_alive=1, timer=test_timer
+    )
     async with AsyncExitStack() as stack:
         await stack.enter_async_context(result_cache)
         async_call_ref = await result_cache.start_call(
@@ -391,7 +393,7 @@ async def test_async_call_pass_auth():
     mock_session = AsyncMock()
     mock_context = Mock()
     mock_context.session = mock_session
-    result_cache = ResultCache(max_size=1, max_keep_alive=1)
+    result_cache = SimpleInMemoryAsyncRequestManager(max_size=1, max_keep_alive=1)
 
     async def test_call(call: types.CallToolRequest) -> types.ServerResult:
         user = user_context.get()