From 20006b3d358116790388288dc0cf83ec462a83c8 Mon Sep 17 00:00:00 2001 From: Ron Mordechai Date: Mon, 7 Jul 2025 23:37:52 +0300 Subject: [PATCH] Extract BaseSession request ID increment logic This allows the request ID increment logic to be overwritten by subclasses. I have a unique need to share sessions' underlying input and output streams, but not the sessions themselves. This would require keeping track of request IDs globally to avoid overlap. Having this logic overridable lets me do that. --- src/mcp/shared/session.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/mcp/shared/session.py b/src/mcp/shared/session.py index b2f49fc8b..427e41e70 100644 --- a/src/mcp/shared/session.py +++ b/src/mcp/shared/session.py @@ -233,8 +233,7 @@ async def send_request( Do not use this method to emit notifications! Use send_notification() instead. """ - request_id = self._request_id - self._request_id = request_id + 1 + request_id = self.next_request_id() response_stream, response_stream_reader = anyio.create_memory_object_stream[JSONRPCResponse | JSONRPCError](1) self._response_streams[request_id] = response_stream @@ -468,3 +467,12 @@ async def _handle_incoming( ) -> None: """A generic handler for incoming messages. Overwritten by subclasses.""" pass + + def next_request_id(self) -> int: + """ + Increment the request ID and return its current value. Can be + overwritten by subclasses. + """ + request_id = self._request_id + self._request_id = request_id + 1 + return request_id