Skip to content

Commit

Permalink
fix: enforce context manager usage for RequestResponder
Browse files Browse the repository at this point in the history
  • Loading branch information
dsp-ant committed Feb 4, 2025
1 parent 08cfbe5 commit b32e455
Showing 1 changed file with 29 additions and 0 deletions.
29 changes: 29 additions & 0 deletions src/mcp/shared/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,21 @@


class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
"""Handles responding to MCP requests and manages request lifecycle.
This class MUST be used as a context manager to ensure proper cleanup and
cancellation handling:
Example:
with request_responder as resp:
await resp.respond(result)
The context manager ensures:
1. Proper cancellation scope setup and cleanup
2. Request completion tracking
3. Cleanup of in-flight requests
"""

def __init__(
self,
request_id: RequestId,
Expand All @@ -55,19 +70,33 @@ def __init__(
self._completed = False
self._cancel_scope = anyio.CancelScope()
self._on_complete = on_complete
self._entered = False # Track if we're in a context manager

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

def __exit__(self, exc_type, exc_val, exc_tb) -> None:
"""Exit the context manager, performing cleanup and notifying completion."""
try:
if self._completed:
self._on_complete(self)
finally:
self._entered = False
self._cancel_scope.__exit__(exc_type, exc_val, exc_tb)

async def respond(self, response: SendResultT | ErrorData) -> None:
"""Send a response for this request.
Must be called within a context manager block.
Raises:
RuntimeError: If not used within a context manager
AssertionError: If request was already responded to
"""
if not self._entered:
raise RuntimeError("RequestResponder must be used as a context manager")
assert not self._completed, "Request already responded to"

if not self.cancelled:
Expand Down

0 comments on commit b32e455

Please sign in to comment.