feat(service): add close() method for graceful connection shutdown #588
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
This PR primarily fixes #572 by enabling graceful shutdown without consuming self. While implementing this, I noticed
delete_session()is spawned as a background task, which meansclose()may return before HTTP session cleanup completes. Since this is part of the same shutdown lifecycle and can cause resource leaks/races, I'm including a small, localized fix to ensure cleanup is completed beforeclose()returns. If maintainers prefer, I can split the cleanup timing change into a follow-up PR.Changes
Service Layer (
service.rs)close(&mut self)- Gracefully shuts down the connection and waits for cleanup to complete without consuming theRunningService. This was the main ask in Client connection is not closed on drop, and no way to gracefully stop #572.close_with_timeout(&mut self, Duration)- Same asclose()but with a bounded wait time. Useful for apps that need deterministic shutdown.is_closed(&self)- Check if the service has already been closed or cancelled.Dropimpl - Logs a debug message if dropped without explicitclose()call. TheDropGuardstill handles async cleanup, but this helps developers catch missing cleanup calls.HTTP Transport (
streamable_http_client.rs)delete_session()from a fire-and-forgettokio::spawnto inline cleanup at the end of the worker'srun()methodclose()from hanging indefinitely if the server is unresponsiveclose()returns, preventing orphaned sessions on authenticated MCP serversWhy both changes together?
The original issue mentions that "for authenticated MCP connections, the sessions on the remote would still be active." The service-layer
close()method alone doesn't fully solve this because the HTTP transport's session deletion was spawned as a background task. By moving it inline with a timeout,close()now guarantees bounded, deterministic cleanup.Design Decisions
Q: Can
close()hang forever?No. The HTTP session cleanup has a 5-second timeout. If the server doesn't respond, we log a warning and continue.
close_with_timeout()provides an additional layer of control at the service level.Q: What happens if cancelled mid-cleanup?
The cleanup runs after the main loop exits, so the cancellation token has already fired. The timeout ensures we don't block indefinitely regardless of server behavior.
Manual Verification
All tests pass, clippy is clean.
Test Plan
test_close_method- Verifiesclose()works and can be called twice safelytest_close_with_timeout- Verifies timeout-bounded shutdowntest_cancel_method- Confirms existingcancel()API still works (backward compat)test_drop_without_close- Verifies drop behavior doesn't panic and async cleanup still happensBackward Compatibility
cancel()method still works (now delegates toclose()internally)waiting()method still worksclose()still triggers cleanup viaDropGuard- just with a debug log nowFixes #572