feat: add request cancellation and in-flight request tracking

This commit adds support for request cancellation and tracking of
in-flight requests in the MCP protocol implementation. The key
architectural changes are:

1. Request Lifecycle Management:
   - Added _in_flight dictionary to BaseSession to track active requests
   - Requests are tracked from receipt until completion/cancellation
   - Added proper cleanup via on_complete callback

2. Cancellation Support:
   - Added CancelledNotification handling in _receive_loop
   - Implemented cancel() method in RequestResponder
   - Uses anyio.CancelScope for robust cancellation
   - Sends error response on cancellation

3. Request Context:
   - Added request_ctx ContextVar for request context
   - Ensures proper cleanup after request handling
   - Maintains request state throughout lifecycle

4. Error Handling:
   - Improved error propagation for cancelled requests
   - Added proper cleanup of cancelled requests
   - Maintains consistency of in-flight tracking

This change enables clients to cancel long-running requests and
servers to properly clean up resources when requests are cancelled.

Github-Issue:#88
This commit is contained in:
David Soria Parra
2025-01-23 20:10:02 +00:00
parent 888bdd3c34
commit 827e494df4
3 changed files with 190 additions and 13 deletions

View File

@@ -453,10 +453,15 @@ class Server:
logger.debug(f"Received message: {message}")
match message:
case RequestResponder(request=types.ClientRequest(root=req)):
await self._handle_request(
message, req, session, raise_exceptions
)
case (
RequestResponder(
request=types.ClientRequest(root=req)
) as responder
):
with responder:
await self._handle_request(
message, req, session, raise_exceptions
)
case types.ClientNotification(root=notify):
await self._handle_notification(notify)