StreamableHttp - Server transport with state management (#553)

src/mcp/shared/session.py

@@ -6,7 +6,6 @@ from types import TracebackType
 from typing import Any, Generic, TypeVar
 
 import anyio
-import anyio.lowlevel
 import httpx
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from pydantic import BaseModel
@@ -24,6 +23,7 @@ from mcp.types import (
     JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
+    NotificationParams,
     RequestParams,
     ServerNotification,
     ServerRequest,
@@ -274,16 +274,32 @@ class BaseSession(
             await response_stream.aclose()
             await response_stream_reader.aclose()
 
-    async def send_notification(self, notification: SendNotificationT) -> None:
+    async def send_notification(
+        self,
+        notification: SendNotificationT,
+        related_request_id: RequestId | None = None,
+    ) -> None:
         """
         Emits a notification, which is a one-way message that does not expect
         a response.
         """
+        # Some transport implementations may need to set the related_request_id
+        # to attribute to the notifications to the request that triggered them.
+        if related_request_id is not None and notification.root.params is not None:
+            # Create meta if it doesn't exist
+            if notification.root.params.meta is None:
+                meta_dict = {"related_request_id": related_request_id}
+
+            else:
+                meta_dict = notification.root.params.meta.model_dump(
+                    by_alias=True, mode="json", exclude_none=True
+                )
+                meta_dict["related_request_id"] = related_request_id
+            notification.root.params.meta = NotificationParams.Meta(**meta_dict)
         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))
 
     async def _send_response(