Skip to content

Commit b32e455

Browse files
committed
fix: enforce context manager usage for RequestResponder
1 parent 08cfbe5 commit b32e455

File tree

1 file changed

+29
-0
lines changed

1 file changed

+29
-0
lines changed

src/mcp/shared/session.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,21 @@
4040

4141

4242
class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
43+
"""Handles responding to MCP requests and manages request lifecycle.
44+
45+
This class MUST be used as a context manager to ensure proper cleanup and
46+
cancellation handling:
47+
48+
Example:
49+
with request_responder as resp:
50+
await resp.respond(result)
51+
52+
The context manager ensures:
53+
1. Proper cancellation scope setup and cleanup
54+
2. Request completion tracking
55+
3. Cleanup of in-flight requests
56+
"""
57+
4358
def __init__(
4459
self,
4560
request_id: RequestId,
@@ -55,19 +70,33 @@ def __init__(
5570
self._completed = False
5671
self._cancel_scope = anyio.CancelScope()
5772
self._on_complete = on_complete
73+
self._entered = False # Track if we're in a context manager
5874

5975
def __enter__(self) -> "RequestResponder[ReceiveRequestT, SendResultT]":
76+
"""Enter the context manager, enabling request cancellation tracking."""
77+
self._entered = True
6078
self._cancel_scope.__enter__()
6179
return self
6280

6381
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
82+
"""Exit the context manager, performing cleanup and notifying completion."""
6483
try:
6584
if self._completed:
6685
self._on_complete(self)
6786
finally:
87+
self._entered = False
6888
self._cancel_scope.__exit__(exc_type, exc_val, exc_tb)
6989

7090
async def respond(self, response: SendResultT | ErrorData) -> None:
91+
"""Send a response for this request.
92+
93+
Must be called within a context manager block.
94+
Raises:
95+
RuntimeError: If not used within a context manager
96+
AssertionError: If request was already responded to
97+
"""
98+
if not self._entered:
99+
raise RuntimeError("RequestResponder must be used as a context manager")
71100
assert not self._completed, "Request already responded to"
72101

73102
if not self.cancelled:

0 commit comments

Comments
 (0)