Skip to content

Complete batch request tasks when withBatch fails - #328

Open
bbdu3 wants to merge 2 commits into
modelcontextprotocol:mainfrom
bbdu3:codex/fix-batch-failure-cleanup
Open

bbdu3 wants to merge 2 commits into
modelcontextprotocol:mainfrom
bbdu3:codex/fix-batch-failure-cleanup

Conversation

@bbdu3

@bbdu3 bbdu3 commented Nov 9, 2026

Copy link
Copy Markdown

Summary

When Client.withBatch throws, tasks returned by Batch.addRequest can remain suspended. Registration happens asynchronously, so cleaning up only the requests already registered at failure time misses continuations that register later.

This change coordinates batch completion and request registration on the Client actor. It handles both transport send failures and a body closure that throws after adding requests.

Fixes #324.

Changes

  • Resume pending requests in the failed batch with the original error.
  • Retain the failure until all expected continuations have registered, allowing late registrations to receive the same error.
  • Avoid resuming continuations already completed by response handling, cancellation, or disconnect.
  • Release batch tracking state once registration completes.

Tests

Validated the final implementation in an isolated WSL Ubuntu 22.04 environment with Swift 6.1:

  • swift build passed.
  • All 16 ClientTests passed.
  • The full Linux test suite passed: 500 tests.
  • git diff --check passed.

Regression coverage includes send failure, body closure failure, registration before and after failure, multiple late registrations, batch isolation, an already-completed response, and continued connection usability.

As a negative control, the public-API regression test was run against the original implementation. Both send failure and body closure failure reproduced the issue: neither child task completed before disconnect. Both cases pass with the fix.

Signed-off-by: bbdu3 <ergouyang854@gamil.com>
@bbdu3
bbdu3 requested a review from a team as a code owner November 9, 2026 05:01

@canarik canarik left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @bbdu3, this fixes a real hang. Holding the error in BatchState until every expected ID registers covers late registrations, and neither path resumes a continuation twice. Traced against the old code, the new public-API test times out on the error count, so it does guard the fix.

One thing before merging:

  1. Document the new failure behavior. When withBatch throws, any task from addRequest that hasn't completed now fails with that error, but both doc comments still only say the error is rethrown, so a caller can't tell that awaiting those tasks no longer hangs. Could you add a line to the withBatch and addRequest docs? The rest can be follow-ups.

Smaller, non-blocking:

  • Duplicate request IDs can still leave a task hanging. The spec forbids reusing an ID in a session, but request(id:) lets callers pick one, and withBatch passes the IDs to finishBatch as a Set. If both registrations arrive after a failure, the second creates a fresh state with no error, so its task waits in pendingRequests and the entry leaks. Could addRequest throw on a repeated ID, with a test? Counting occurrences wouldn't help, since pendingRequests holds one continuation per ID.
  • On the success path, registration still happens in a separate Task, so a response that handleBatchResponse sees first is dropped, and a registration after disconnect() is left waiting until a later disconnect. That predates this PR, so a follow-up is fine.
  • addPendingBatchRequest and finishBatch are internal only for the tests; a short comment would help, since addPendingRequest next to them is private.
  • Each registration copies BatchState out of the dictionary, so the set is copied every time and a large batch registers in quadratic time; mutating the entry in place avoids that.

Signed-off-by: bbdu3 <ergouyang854@gamil.com>
@bbdu3

bbdu3 commented Nov 10, 2026

Copy link
Copy Markdown
Author

@canarik ,thanks for review, I’ve updated the withBatch and addRequest doc comments to explain that unfinished result tasks fail with the same error when withBatch throws, and synchronized the README.

I also addressed the smaller non-blocking items: rejecting duplicate IDs within a batch with regression tests, mutating BatchState in place to avoid repeated set copies, and documenting why the two helpers remain internal for testing.

I’ve left the success-path and post-disconnect registration races for a follow-up. Those require broader changes to how registration coordinates with sending, response handling, and disconnect, plus dedicated concurrency tests. Keeping them separate preserves the scope of this failure-cleanup fix. Duplicate-ID validation is also limited to the current batch; it does not track IDs across the entire session.

The build and all 502 Linux tests passed with Swift 6.1, and DocC passed with --warnings-as-errors. The new duplicate-ID tests reproduce continuation leaks against the previous implementation.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Client.withBatch leaves per-request tasks pending when transport.send throws

2 participants