Add strict mode to pyright (#315)

* Add strict mode to pyright

* Apply UP rule

* fix readme

* More correct

* Leave wrong Context for now

* Add strict mode to pyright

* Apply UP rule

* fix readme

* fix

* ignore

src/mcp/shared/memory.py

@@ -2,9 +2,10 @@
 In-memory transports
 """
 
+from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
 from datetime import timedelta
-from typing import AsyncGenerator
+from typing import Any
 
 import anyio
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
@@ -52,7 +53,7 @@ async def create_client_server_memory_streams() -> (
 
 @asynccontextmanager
 async def create_connected_server_and_client_session(
-    server: Server,
+    server: Server[Any],
     read_timeout_seconds: timedelta | None = None,
     sampling_callback: SamplingFnT | None = None,
     list_roots_callback: ListRootsFnT | None = None,

src/mcp/shared/progress.py

@@ -1,10 +1,19 @@
+from collections.abc import Generator
 from contextlib import contextmanager
 from dataclasses import dataclass, field
+from typing import Generic
 
 from pydantic import BaseModel
 
 from mcp.shared.context import RequestContext
-from mcp.shared.session import BaseSession
+from mcp.shared.session import (
+    BaseSession,
+    ReceiveNotificationT,
+    ReceiveRequestT,
+    SendNotificationT,
+    SendRequestT,
+    SendResultT,
+)
 from mcp.types import ProgressToken
 
 
@@ -14,8 +23,22 @@ class Progress(BaseModel):
 
 
 @dataclass
-class ProgressContext:
-    session: BaseSession
+class ProgressContext(
+    Generic[
+        SendRequestT,
+        SendNotificationT,
+        SendResultT,
+        ReceiveRequestT,
+        ReceiveNotificationT,
+    ]
+):
+    session: BaseSession[
+        SendRequestT,
+        SendNotificationT,
+        SendResultT,
+        ReceiveRequestT,
+        ReceiveNotificationT,
+    ]
     progress_token: ProgressToken
     total: float | None
     current: float = field(default=0.0, init=False)
@@ -29,7 +52,27 @@ class ProgressContext:
 
 
 @contextmanager
-def progress(ctx: RequestContext, total: float | None = None):
+def progress(
+    ctx: RequestContext[
+        BaseSession[
+            SendRequestT,
+            SendNotificationT,
+            SendResultT,
+            ReceiveRequestT,
+            ReceiveNotificationT,
+        ]
+    ],
+    total: float | None = None,
+) -> Generator[
+    ProgressContext[
+        SendRequestT,
+        SendNotificationT,
+        SendResultT,
+        ReceiveRequestT,
+        ReceiveNotificationT,
+    ],
+    None,
+]:
     if ctx.meta is None or ctx.meta.progressToken is None:
         raise ValueError("No progress token provided")
 

src/mcp/shared/session.py

@@ -1,7 +1,9 @@
 import logging
+from collections.abc import Callable
 from contextlib import AsyncExitStack
 from datetime import timedelta
-from typing import Any, Callable, Generic, TypeVar
+from types import TracebackType
+from typing import Any, Generic, TypeVar
 
 import anyio
 import anyio.lowlevel
@@ -86,7 +88,12 @@ class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
         self._cancel_scope.__enter__()
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
         """Exit the context manager, performing cleanup and notifying completion."""
         try:
             if self._completed:
@@ -112,7 +119,7 @@ class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
         if not self.cancelled:
             self._completed = True
 
-            await self._session._send_response(
+            await self._session._send_response(  # type: ignore[reportPrivateUsage]
                 request_id=self.request_id, response=response
             )
 
@@ -126,7 +133,7 @@ class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
         self._cancel_scope.cancel()
         self._completed = True  # Mark as completed so it's removed from in_flight
         # Send an error response to indicate cancellation
-        await self._session._send_response(
+        await self._session._send_response(  # type: ignore[reportPrivateUsage]
             request_id=self.request_id,
             response=ErrorData(code=0, message="Request cancelled", data=None),
         )
@@ -137,7 +144,7 @@ class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
 
     @property
     def cancelled(self) -> bool:
-        return self._cancel_scope is not None and self._cancel_scope.cancel_called
+        return self._cancel_scope.cancel_called
 
 
 class BaseSession(
@@ -202,7 +209,12 @@ class BaseSession(
         self._task_group.start_soon(self._receive_loop)
         return self
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool | None:
         await self._exit_stack.aclose()
         # Using BaseSession as a context manager should not block on exit (this
         # would be very surprising behavior), so make sure to cancel the tasks
@@ -324,7 +336,7 @@ class BaseSession(
 
                     self._in_flight[responder.request_id] = responder
                     await self._received_request(responder)
-                    if not responder._completed:
+                    if not responder._completed:  # type: ignore[reportPrivateUsage]
                         await self._incoming_message_stream_writer.send(responder)
 
                 elif isinstance(message.root, JSONRPCNotification):