add test for keep_alive expiry

Author: davemssavage

src/mcp/server/lowlevel/result_cache.py

@@ -1,8 +1,8 @@
+import time
 from collections.abc import Awaitable, Callable
 from concurrent.futures import Future
 from dataclasses import dataclass, field
 from logging import getLogger
-from time import time
 from types import TracebackType
 from typing import Any
 from uuid import uuid4
@@ -23,12 +23,21 @@
 @dataclass
 class InProgress:
     token: str
+    timer: Callable[[], float]
     user: AuthenticatedUser | None = None
     future: Future[types.CallToolResult] | None = None
     sessions: dict[int, ServerSession] = field(default_factory=lambda: {})
     session_progress: dict[int, types.ProgressToken | None] = field(
         default_factory=lambda: {}
     )
+    keep_alive: int | None = None
+    keep_alive_start: int | None = None
+
+    def is_expired(self):
+        if self.keep_alive_start is None or self.keep_alive is None:
+            return False
+        else:
+            return int(self.timer()) > self.keep_alive_start + self.keep_alive
 
 
 class ResultCache:
@@ -54,11 +63,17 @@ class ResultCache:
     _session_lookup: dict[int, types.AsyncToken]
     _portal: BlockingPortal
 
-    def __init__(self, max_size: int, max_keep_alive: int):
+    def __init__(
+        self,
+        max_size: int,
+        max_keep_alive: int,
+        timer: Callable[[], float] = time.monotonic,
+    ):
         self._max_size = max_size
         self._max_keep_alive = max_keep_alive
         self._in_progress = {}
         self._session_lookup = {}
+        self._timer = timer
         self._portal_provider = BlockingPortalProvider()
 
     async def __aenter__(self):
@@ -105,6 +120,7 @@ async def call_tool():
         in_progress.user = user_context.get()
         session_id = id(ctx.session)
         in_progress.sessions[session_id] = ctx.session
+        in_progress.keep_alive = timeout
         if req.params.meta is not None:
             progress_token = req.params.meta.progressToken
         else:
@@ -114,7 +130,7 @@ async def call_tool():
         in_progress.future = self._portal.start_task_soon(call_tool)
         result = types.CallToolAsyncResult(
             token=in_progress.token,
-            recieved=round(time()),
+            recieved=round(self._timer()),
             keepAlive=timeout,
             accepted=True,
         )
@@ -176,6 +192,15 @@ async def get_result(self, req: types.GetToolAsyncResultRequest):
             if in_progress.user == user_context.get():
                 assert in_progress.future is not None
                 # TODO add timeout to get async result
+                if in_progress.is_expired():
+                    self._portal.start_task_soon(self._expire)
+                    return types.CallToolResult(
+                        content=[
+                            types.TextContent(type="text", text="Unknown async token")
+                        ],
+                        isError=True,
+                    )
+
                 try:
                     result = in_progress.future.result(1)
                     logger.debug(f"Found result {result}")
@@ -235,7 +260,12 @@ async def session_close_hook(self, session: ServerSession):
             if found is None:
                 logger.warning("No session found")
             if len(in_progress.sessions) == 0:
-                self._in_progress.pop(dropped, None)
+                in_progress.keep_alive_start = int(self._timer())
+
+    async def _expire(self):
+        for in_progress in self._in_progress.values():
+            if in_progress.is_expired():
+                self._in_progress.pop(in_progress.token, None)
                 assert in_progress.future is not None
                 logger.debug("Cancelled in progress future")
                 in_progress.future.cancel()
@@ -251,6 +281,6 @@ async def _new_in_progress(self) -> InProgress:
             # for context
             token = str(uuid4())
             if token not in self._in_progress:
-                new_in_progress = InProgress(token)
+                new_in_progress = InProgress(token, self._timer)
                 self._in_progress[token] = new_in_progress
                 return new_in_progress

tests/server/lowlevel/test_result_cache.py

@@ -68,7 +68,6 @@ async def test_call(call: types.CallToolRequest) -> types.ServerResult:
     mock_context_2 = Mock()
 
     mock_context_2.session = mock_session_2
-    mock_session_2.send_progress_notification.result = None
 
     result_cache = ResultCache(max_size=1, max_keep_alive=1)
     async with AsyncExitStack() as stack:
@@ -123,3 +122,178 @@ async def test_call(call: types.CallToolRequest) -> types.ServerResult:
             message=None,
             resource_uri=None,
         )
+
+
+@pytest.mark.anyio
+async def test_async_call_keep_alive():
+    """Tests async call keep alive"""
+
+    async def test_call(call: types.CallToolRequest) -> types.ServerResult:
+        return types.ServerResult(
+            types.CallToolResult(content=[types.TextContent(type="text", text="test")])
+        )
+
+    async_call = types.CallToolAsyncRequest(
+        method="tools/async/call", params=types.CallToolAsyncRequestParams(name="test")
+    )
+
+    mock_session_1 = AsyncMock()
+    mock_context_1 = Mock()
+    mock_context_1.session = mock_session_1
+
+    mock_session_2 = AsyncMock()
+    mock_context_2 = Mock()
+
+    mock_context_2.session = mock_session_2
+
+    result_cache = ResultCache(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(
+            test_call, async_call, mock_context_1
+        )
+        assert async_call_ref.token is not None
+
+        await result_cache.session_close_hook(mock_session_1)
+
+        await result_cache.join_call(
+            req=types.JoinCallToolAsyncRequest(
+                method="tools/async/join",
+                params=types.JoinCallToolRequestParams(
+                    token=async_call_ref.token,
+                    _meta=types.RequestParams.Meta(progressToken="test"),
+                ),
+            ),
+            ctx=mock_context_2,
+        )
+        assert async_call_ref.token is not None
+        await result_cache.notification_hook(
+            session=mock_session_1,
+            notification=types.ServerNotification(
+                types.ProgressNotification(
+                    method="notifications/progress",
+                    params=types.ProgressNotificationParams(
+                        progressToken="test", progress=1
+                    ),
+                )
+            ),
+        )
+
+        result = await result_cache.get_result(
+            types.GetToolAsyncResultRequest(
+                method="tools/async/get",
+                params=types.GetToolAsyncResultRequestParams(
+                    token=async_call_ref.token
+                ),
+            )
+        )
+
+        assert not result.isError, str(result)
+        assert not result.isPending
+        assert len(result.content) == 1
+        assert type(result.content[0]) is types.TextContent
+        assert result.content[0].text == "test"
+
+
+@pytest.mark.anyio
+async def test_async_call_keep_alive_expired():
+    """Tests async call keep alive expiry"""
+
+    async def test_call(call: types.CallToolRequest) -> types.ServerResult:
+        return types.ServerResult(
+            types.CallToolResult(content=[types.TextContent(type="text", text="test")])
+        )
+
+    async_call = types.CallToolAsyncRequest(
+        method="tools/async/call", params=types.CallToolAsyncRequestParams(name="test")
+    )
+
+    mock_session_1 = AsyncMock()
+    mock_context_1 = Mock()
+    mock_context_1.session = mock_session_1
+
+    mock_session_2 = AsyncMock()
+    mock_context_2 = Mock()
+    mock_context_2.session = mock_session_2
+
+    mock_session_3 = AsyncMock()
+    mock_context_3 = Mock()
+    mock_context_3.session = mock_session_3
+
+    time = 0.0
+
+    def test_timer():
+        return time
+
+    result_cache = ResultCache(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(
+            test_call, async_call, mock_context_1
+        )
+        assert async_call_ref.token is not None
+
+        # lose the connection
+        await result_cache.session_close_hook(mock_session_1)
+
+        # reconnect before keep_alive_timeout
+        time = 0.5
+        await result_cache.join_call(
+            req=types.JoinCallToolAsyncRequest(
+                method="tools/async/join",
+                params=types.JoinCallToolRequestParams(
+                    token=async_call_ref.token,
+                    _meta=types.RequestParams.Meta(progressToken="test"),
+                ),
+            ),
+            ctx=mock_context_2,
+        )
+
+        result = await result_cache.get_result(
+            types.GetToolAsyncResultRequest(
+                method="tools/async/get",
+                params=types.GetToolAsyncResultRequestParams(
+                    token=async_call_ref.token
+                ),
+            )
+        )
+
+        # should successfully read data
+        assert not result.isError, str(result)
+        assert len(result.content) == 1
+        assert type(result.content[0]) is types.TextContent
+        assert result.content[0].text == "test"
+
+        # lose connection a second time
+
+        await result_cache.session_close_hook(mock_session_2)
+
+        time = 2
+
+        # reconnect after the keep_alive_timeout
+
+        await result_cache.join_call(
+            req=types.JoinCallToolAsyncRequest(
+                method="tools/async/join",
+                params=types.JoinCallToolRequestParams(
+                    token=async_call_ref.token,
+                    _meta=types.RequestParams.Meta(progressToken="test"),
+                ),
+            ),
+            ctx=mock_context_3,
+        )
+
+        result = await result_cache.get_result(
+            types.GetToolAsyncResultRequest(
+                method="tools/async/get",
+                params=types.GetToolAsyncResultRequestParams(
+                    token=async_call_ref.token
+                ),
+            )
+        )
+
+        # now token should be expired
+        assert result.isError, str(result)
+        assert len(result.content) == 1
+        assert type(result.content[0]) is types.TextContent
+        assert result.content[0].text == "Unknown async token"