fix: improve session cleanup and progress handling

1. Add proper session cleanup handling
   - Track session state with _closed flag
   - Handle cancellation and cleanup errors gracefully
   - Skip notification validation during cleanup

2. Improve progress context
   - Add final_progress method
   - Send completion notification in finally block
   - Handle progress cleanup properly

This fix addresses issues with session cleanup causing validation
errors and improves progress notification reliability.

Author: donghao1393

src/mcp/shared/progress.py

@@ -1,4 +1,4 @@
-from contextlib import contextmanager
+from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 
 from pydantic import BaseModel
@@ -21,20 +21,34 @@ class ProgressContext:
     current: float = field(default=0.0, init=False)
 
     async def progress(self, amount: float) -> None:
+        """Update progress by the given amount and send notification."""
         self.current += amount
-
         await self.session.send_progress_notification(
             self.progress_token, self.current, total=self.total
         )
 
+    async def final_progress(self) -> None:
+        """Send the final progress notification."""
+        if self.total is not None and self.current < self.total:
+            self.current = self.total
+            await self.session.send_progress_notification(
+                self.progress_token, self.current, total=self.total
+            )
+
+
+@asynccontextmanager
+async def progress(ctx: RequestContext, total: float | None = None):
+    """Context manager for progress tracking and notification.
 
-@contextmanager
-def progress(ctx: RequestContext, total: float | None = None):
+    Args:
+        ctx: Request context containing the session and progress token
+        total: Optional total progress amount
+    """
     if ctx.meta is None or ctx.meta.progressToken is None:
         raise ValueError("No progress token provided")
 
     progress_ctx = ProgressContext(ctx.session, ctx.meta.progressToken, total)
     try:
         yield progress_ctx
     finally:
-        pass
+        await progress_ctx.final_progress()

src/mcp/shared/session.py

@@ -193,6 +193,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
         # Using BaseSession as a context manager should not block on exit (this
         # would be very surprising behavior), so make sure to cancel the tasks
         # in the task group.
+        self._closed = True
         self._task_group.cancel_scope.cancel()
         return await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
 
@@ -256,12 +257,21 @@ async def send_notification(self, notification: SendNotificationT) -> None:
         Emits a notification, which is a one-way message that does not expect
         a response.
         """
-        jsonrpc_notification = JSONRPCNotification(
-            jsonrpc="2.0",
-            **notification.model_dump(by_alias=True, mode="json", exclude_none=True),
-        )
+        # Skip sending notifications if the session is closed
+        if self._closed:
+            return
+
+        try:
+            jsonrpc_notification = JSONRPCNotification(
+                jsonrpc="2.0",
+                **notification.model_dump(by_alias=True, mode="json", exclude_none=True),
+            )
 
-        await self._write_stream.send(JSONRPCMessage(jsonrpc_notification))
+            await self._write_stream.send(JSONRPCMessage(jsonrpc_notification))
+        except Exception:
+            # Ignore notification send errors during session cleanup
+            if not self._closed:
+                raise
 
     async def _send_response(
         self, request_id: RequestId, response: SendResultT | ErrorData
@@ -279,6 +289,19 @@ async def _send_response(
             )
             await self._write_stream.send(JSONRPCMessage(jsonrpc_response))
 
+    def _should_validate_notification(self, message_root: JSONRPCNotification) -> bool:
+        """
+        Determines if a notification should be validated.
+        Internal notifications (like cancelled) should be ignored.
+        """
+        try:
+            return (
+                getattr(message_root, "method", None) != "cancelled" and
+                not self._closed
+            )
+        except:
+            return False
+
     async def _receive_loop(self) -> None:
         async with (
             self._read_stream,
@@ -378,4 +401,4 @@ def incoming_messages(
         | ReceiveNotificationT
         | Exception
     ]:
-        return self._incoming_message_stream_reader
+        return self._incoming_message_stream_reader