formatting and other misc tidy up highlighted by ruff check

davemssavage · davemssavage · commit ba0f9ee466a9 · 2025-06-03T05:47:50.000Z
diff --git a/src/mcp/server/lowlevel/result_cache.py b/src/mcp/server/lowlevel/result_cache.py
@@ -26,7 +26,10 @@ class InProgress:
     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: {})
+    session_progress: dict[int, types.ProgressToken | None] = field(
+        default_factory=lambda: {}
+    )
+
 
 class ResultCache:
     """
@@ -100,9 +103,14 @@ async def call_tool():
             return result.root
 
         in_progress.user = user_context.get()
-        in_progress.sessions[id(ctx.session)] = ctx.session
-        in_progress.session_progress[id(ctx.session)] = None if req.params.meta is None else req.params.meta.progressToken
-        self._session_lookup[id(ctx.session)] = in_progress.token
+        session_id = id(ctx.session)
+        in_progress.sessions[session_id] = ctx.session
+        if req.params.meta is not None:
+            progress_token = req.params.meta.progressToken
+        else:
+            progress_token = None
+        in_progress.session_progress[session_id] = progress_token
+        self._session_lookup[session_id] = in_progress.token
         in_progress.future = self._portal.start_task_soon(call_tool)
         result = types.CallToolAsyncResult(
             token=in_progress.token,
@@ -127,10 +135,15 @@ async def join_call(
         else:
             # TODO consider adding authorisation layer to make this decision
             if in_progress.user == user_context.get():
-                logger.debug(f"Received join from {id(ctx.session)}")
-                self._session_lookup[id(ctx.session)] = req.params.token
-                in_progress.sessions[id(ctx.session)] = ctx.session
-                in_progress.session_progress[id(ctx.session)] = None if req.params.meta is None else req.params.meta.progressToken
+                session_id = id(ctx.session)
+                logger.debug(f"Received join from {session_id}")
+                self._session_lookup[session_id] = req.params.token
+                in_progress.sessions[session_id] = ctx.session
+                if req.params.meta is not None:
+                    progress_token = req.params.meta.progressToken
+                else:
+                    progress_token = None
+                in_progress.session_progress[session_id] = progress_token
                 return types.CallToolAsyncResult(token=req.params.token, accepted=True)
             else:
                 # TODO consider sending error via get result
@@ -196,25 +209,22 @@ async def notification_hook(
                 logger.debug("Discarding progress notification from unknown session")
             else:
                 in_progress = self._in_progress.get(async_token)
-                if in_progress is None:
-                    # this should not happen
-                    logger.error("Discarding progress notification, not async")
-                else:
-                    for session_id, other_session in in_progress.sessions.items():
-                        logger.debug(f"Checking {session_id} == {id(session)}")
-                        if not session_id == id(session):
-                            logger.debug(f"Sending progress to {id(other_session)}")
-                            progress_token = in_progress.session_progress.get(id(other_session))
-                            assert progress_token is not None
-                            await other_session.send_progress_notification(
-                                # TODO this token is incorrect 
-                                # it needs to be collected from original request
-                                progress_token=progress_token,
-                                progress=notification.root.params.progress,
-                                total=notification.root.params.total,
-                                message=notification.root.params.message,
-                                resource_uri=notification.root.params.resourceUri,
-                            )
+                assert in_progress is not None
+                for other_id, other_session in in_progress.sessions.items():
+                    logger.debug(f"Checking {other_id} == {id(session)}")
+                    if not other_id == id(session):
+                        logger.debug(f"Sending progress to {other_id}")
+                        progress_token = in_progress.session_progress.get(other_id)
+                        assert progress_token is not None
+                        await other_session.send_progress_notification(
+                            # TODO this token is incorrect
+                            # it needs to be collected from original request
+                            progress_token=progress_token,
+                            progress=notification.root.params.progress,
+                            total=notification.root.params.total,
+                            message=notification.root.params.message,
+                            resource_uri=notification.root.params.resourceUri,
+                        )
 
     async def session_close_hook(self, session: ServerSession):
         logger.debug(f"Closing {id(session)}")
diff --git a/tests/server/lowlevel/test_result_cache.py b/tests/server/lowlevel/test_result_cache.py
@@ -1,24 +1,23 @@
+from contextlib import AsyncExitStack
+from unittest.mock import AsyncMock, Mock
+
 import pytest
+
 from mcp import types
 from mcp.server.lowlevel.result_cache import ResultCache
-from unittest.mock import AsyncMock, Mock, patch
-from contextlib import AsyncExitStack
+
 
 @pytest.mark.anyio
 async def test_async_call():
     """Tests basic async call"""
+
     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"
+        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 = AsyncMock()
@@ -27,37 +26,38 @@ async def test_call(call: types.CallToolRequest) -> types.ServerResult:
     result_cache = ResultCache(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(test_call, async_call, mock_context)
+        async_call_ref = await result_cache.start_call(
+            test_call, async_call, mock_context
+        )
         assert async_call_ref.token is not None
 
-        result = await result_cache.get_result(types.GetToolAsyncResultRequest(
-            method="tools/async/get",
-            params=types.GetToolAsyncResultRequestParams(
-                token = async_call_ref.token
+        result = await result_cache.get_result(
+            types.GetToolAsyncResultRequest(
+                method="tools/async/get",
+                params=types.GetToolAsyncResultRequestParams(
+                    token=async_call_ref.token
+                ),
             )
-        ))
+        )
 
         assert not result.isError
         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_join_call_progress():
     """Tests basic async call"""
+
     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"
+        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()
@@ -73,38 +73,42 @@ async def test_call(call: types.CallToolRequest) -> types.ServerResult:
     result_cache = ResultCache(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(test_call, async_call, mock_context_1)
+        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.join_call(
             req=types.JoinCallToolAsyncRequest(
                 method="tools/async/join",
                 params=types.JoinCallToolRequestParams(
                     token=async_call_ref.token,
-                    _meta = types.RequestParams.Meta(
-                        progressToken="test"
-                    )
-                )
+                    _meta=types.RequestParams.Meta(progressToken="test"),
+                ),
             ),
-            ctx=mock_context_2
+            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
+            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
+        result = await result_cache.get_result(
+            types.GetToolAsyncResultRequest(
+                method="tools/async/get",
+                params=types.GetToolAsyncResultRequestParams(
+                    token=async_call_ref.token
+                ),
             )
-        ))
+        )
 
         assert not result.isError
         assert not result.isPending
@@ -113,9 +117,9 @@ async def test_call(call: types.CallToolRequest) -> types.ServerResult:
         assert result.content[0].text == "test"
         mock_context_1.send_progress_notification.assert_not_called()
         mock_session_2.send_progress_notification.assert_called_with(
-            progress_token="test", 
-            progress=1.0, 
-            total=None, 
-            message=None, 
-            resource_uri = None
+            progress_token="test",
+            progress=1.0,
+            total=None,
+            message=None,
+            resource_uri=None,
         )
diff --git a/tests/shared/test_session.py b/tests/shared/test_session.py
@@ -1,4 +1,5 @@
 from collections.abc import AsyncGenerator
+from logging import getLogger
 
 import anyio
 import pytest
@@ -19,6 +20,8 @@
     EmptyResult,
 )
 
+logger = getLogger(__name__)
+
 
 @pytest.fixture
 def mcp_server() -> Server:
@@ -211,15 +214,14 @@ async def get_result(client_session: ClientSession, async_token: types.AsyncToke
         assert type(result.content[0]) is types.TextContent
         assert result.content[0].text == "test"
 
-from logging import getLogger
-
-logger = getLogger(__name__)
 
 @pytest.mark.anyio
-@pytest.mark.skip(reason="This test does not work, there is a subtle "
-                  "bug with event.wait, lower level test_result_cache "
-                  "tests underlying behaviour, revisit with feedback " \
-                  "from someone who cah help debug")
+@pytest.mark.skip(
+    reason="This test does not work, there is a subtle "
+    "bug with event.wait, lower level test_result_cache "
+    "tests underlying behaviour, revisit with feedback "
+    "from someone who cah help debug"
+)
 async def test_request_async_join():
     """Test that requests can be joined from external sessions."""
     # The tool is already registered in the fixture
@@ -236,7 +238,6 @@ async def test_request_async_join():
     ev_client_2_progressed_1 = anyio.Event()
     ev_done = anyio.Event()
 
-
     # Start the request in a separate task so we can cancel it
     def make_server() -> Server:
         server = Server(name="TestSessionServer")
@@ -328,8 +329,7 @@ async def client_2_progress_callback(
         logger.info(f"client2: progress done: {progress}/{total}")
 
     async def join_request(
-        client_session: ClientSession, 
-        async_token: types.AsyncToken
+        client_session: ClientSession, async_token: types.AsyncToken
     ):
         return await client_session.send_request(
             ClientRequest(
@@ -365,6 +365,7 @@ async def get_result(client_session: ClientSession, async_token: types.AsyncToke
     token = None
 
     async with anyio.create_task_group() as tg:
+
         async def client_1_submit():
             async with create_connected_server_and_client_session(
                 server
@@ -430,6 +431,7 @@ async def client_2_join():
     assert ev_client_1_progressed_2.is_set()
     assert ev_client_2_progressed_1.is_set()
 
+
 @pytest.mark.anyio
 async def test_connection_closed():
     """
