Fix ATTENTION resynchronization during row reads

[Saurabh Singh (saurabh500)]

Description

Cancellation and request timeouts could drop a TDS parser while a ROW or NBCROW token was only partially consumed. The following ATTENTION drain then lacked the parser state and column metadata needed to find the next token boundary, so the session could not be safely reused. Async Python cancellation also dropped the Rust fetch future before protocol settlement completed.

This change keeps the in-flight parser alive while ATTENTION is sent, completes the interrupted token under the same bounded deadline, and drains rows and control tokens with the active or newly received COLMETADATA through DONE_ATTN. The resulting control-token state is applied to TdsClient, which then clears the cancelled result state while retaining a synchronized connection. Python fetch and nextset work now continues in a detached Tokio task until settlement, with operation ownership held across cancellation and the late-cancellation race handled explicitly. If nextset advances to a row result whose Python description cannot be materialized, the cursor retains ownership so a later nextset can drain the unread rows before advancing.

Validation includes workspace and Python formatting/clippy, 3,303 workspace unit tests, 226 Python-core Rust tests, 15 targeted ATTENTION tests, and 6 mock-backed Python fetch tests. The live SQL cancellation tests were added but could not run locally because the available SQL Server rejected the configured test credentials; the full cargo btest integration surface is therefore left unchecked for CI.

Related Issues

Closes #471

Checklist

• cargo bfmt passes
• cargo bclippy passes
• cargo btest passes
• New/changed functionality has tests
• Public API changes are documented (N/A: no public API changes)

[Saurabh Singh (saurabh500)]

Summary

This replaces drop-the-parser cancellation with a design that keeps the in-flight parser future alive across ATTENTION: SharedStream lets the ATTENTION packet be written while the parser still borrows the transport, await_read_or_interrupt hands back Pin<&mut F> instead of dropping it, and the drain now decodes ROW/NBCROW with real column metadata before accepting DONE_ATTN. The mechanism is sound and the client-side settlement is wired at every interruptible read site. No blocking findings. The one that matters most is a new test that stays green when you break the mechanism its name claims to guard.

Verified against the code, not assumed:

• settle_interrupted_read covers every async/interruptible self.transport read in tds_client.rs. The five call sites it does not cover (5811, 6088, 6151, 6200, 6238) are synchronous try_* buffered helpers that can neither time out nor be cancelled, so they need no settlement.
• The Python late-cancellation path only sends ATTENTION if client.has_open_batch(), so normalize_after_attention() — which sets self.cancel_handle = None — can be skipped after cancel_fetch already cancelled the root handle. That is safe because tds_client.rs:1449 re-arms self.cancel_handle from the caller's handle at the start of every command, so a cancelled token cannot survive into the next operation.
• Widening ATTN acceptance to Done / DoneProc / DoneInProc matches SqlClient, where TryProcessDone is shared across all three token types.
• The send_attention_packet write has no explicit flush, which is parity with the existing NetworkWriter::send for NetworkTransport — not a defect.
• The hand-rolled biased select produces the same OperationCancelledError("Request was cancelled") string as CancelHandle::run_until_cancelled, so no caller sees a changed error.
• The run_until_cancelled read wrappers still in io/token_stream.rs are #[cfg(fuzzing)], not a production sibling left unconverted.

What I ran, and what I did not. cargo nextest run -p mssql-tds --lib --no-fail-fast: 2116 tests, 2109 pass. The 7 failures are certificate_validator and win_tls::validate looking for tests/test_certificates/*.pem — the known missing-fixture set on a clean tree, unrelated to this change. I did not run any live-SQL suite. When I wrote this, the ADO Windows/Linux/macOS build+test jobs and coverage-report were still pending on e21247a, so the unchecked cargo btest box in the checklist is still unconfirmed by CI. The new Python @pytest.mark.integration tests do run in CI — dev/test-python.sh:113 selects -m "not longhaul and not smoke" when a live server is configured — so they will be exercised once that Linux job finishes.

Findings

Blocking — none.

Suggestion — four, left inline: a test that does not pin its own mechanism (network_transport.rs:5947), a now-dead request_timeout module (transport.rs:17), two divergent ways to send ATTENTION (network_transport.rs:723), and an unadvertised ownership change on the nextset materialization-failure path (async_fetch.rs:871).

Nit — two, neither anchorable to a changed line:

• mssql-tds/src/io/token_stream.rs:863 still reads "The fuzzing counterpart of NetworkTransport::cancel_read_stream_and_wait". This PR deletes that function, so the doc points at nothing.
• normalize_after_attention() now ends with set_has_open_batch(false) (tds_client.rs:4137), which makes the explicit set_has_open_batch(false) calls that immediately follow send_attention_with_timeout at tds_client.rs:2032 and tds_client.rs:2559 redundant. Harmless, but they now read as if they were still doing something.

[Copilot]

🟡 Changes recommended

Higher-level drain handlers still retire connections after successfully settled cancellation or timeout errors.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Fixes #471 by preserving parser state during ATTENTION cancellation so synchronized TDS sessions can be reused.

Changes:

• Adds metadata-aware ATTENTION draining and state replay.
• Normalizes TdsClient after interrupted reads.
• Keeps Python fetch operations alive until cancellation settlement completes.

File summaries

File	Description
mssql-tds/src/test_client_support.rs	Updates test transport interface.
mssql-tds/src/io/token_stream.rs	Makes row pause state cloneable.
mssql-tds/src/fuzz_support.rs	Updates fuzz transport interface.
mssql-tds/src/connection/transport/tds_transport.rs	Adds parser context to ATTENTION handling.
mssql-tds/src/connection/transport/request_timeout.rs	Removes superseded timeout helper.
mssql-tds/src/connection/transport/network_transport.rs	Implements parser-preserving ATTENTION settlement.
mssql-tds/src/connection/transport/any_transport.rs	Forwards parser context and settlement state.
mssql-tds/src/connection/transport.rs	Removes obsolete module declaration.
mssql-tds/src/connection/tds_client.rs	Applies settlement tokens and clears cancelled results.
mssql-py-core/tests/test_async_fetch.py	Tests reusable sessions after cancellation.
mssql-py-core/src/async_session.rs	Tracks cancellation while retaining ownership.
mssql-py-core/src/async_fetch.rs	Detaches protocol work until settlement completes.

Review details

• Files reviewed: 12/12 changed files
• Comments generated: 1
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[David Engel (David-Engel)]

Automated review — generated by GitHub Copilot on behalf of David Engel (@David-Engel). This is not an approval and does not satisfy the human review requirement. Findings may be incomplete or wrong; push back on anything that looks off.

Summary

This replaces drop-the-parser cancellation with a design that keeps the interrupted parser future alive across ATTENTION. SharedStream lets the ATTENTION header be written while the parser still borrows NetworkTransport, await_read_or_interrupt hands back Pin<&mut F> instead of dropping it, and drain_to_attention_ack now decodes ROW/NBCROW with real column metadata — seeded from the interrupted context and re-seeded from any COLMETADATA that arrives mid-drain — before accepting DONE_ATTN. The drained control tokens are handed back to TdsClient so transaction/environment/recovery state stays synchronized, and the Python bindings move the TDS work into a detached Tokio task so an asyncio cancellation cannot abandon a parser mid-token. The mechanism reads as sound and I found no blocking issues. Two suggestions below, both in the new client-side settlement helpers.

Verification

Read the full origin/main...HEAD diff (merge base dd96d8e), plus the surrounding unchanged code: every settle_interrupted_read call site in tds_client.rs, apply_drain_side_effect / observe_response_token / retire_after_failed_drain, the finish_command reset block at tds_client.rs:6943-6953, NetworkWriter::send (confirming the missing flush in send_attention_packet is parity, not a defect), PacketWriter::build_header (confirming the hand-built ATTENTION packet is a valid 8-byte 0x06/EOM header), and the claim_execute → claim_fetch → cancel_fetch ownership chain in mssql-py-core/src/async_session.rs.

Ran cargo nextest run -p mssql-tds --lib --no-fail-fast against the assigned target dir: 2111 tests, 2104 pass. The 7 failures are the known missing-fixture set (certificate_validator, win_tls::validate looking for tests/test_certificates/*.pem) and are unrelated to this change.

I also re-ran the mutation that the earlier review used to show queued_rows_are_drained_with_current_metadata was vacuous: forcing attention_drain_context's ParserContext::ColumnMetadata arm to yield metadata: None. With the two-row fixture added in 916620d, the test now fails as intended ("ROW was drained through DONE_ATTN, so the connection is reusable"), while the sibling ATTENTION tests stay green. That earlier finding is genuinely fixed, and the test now pins its own mechanism. I did not run any live-SQL suite; ADO validation and coverage-report were still pending on 2e50c79 when I looked, so the unchecked cargo btest box remains unconfirmed by CI.

Findings

Blocking

None.

Suggestion

Two, left inline in mssql-tds/src/connection/tds_client.rs: normalize_after_attention omits one field that its sibling reset block clears, and consume_attention_settlement can retire a healthy connection over state it is about to discard.

Nit

None.

[David Engel (David-Engel)]

Automated review — generated by GitHub Copilot on behalf of David Engel (@David-Engel). This is not an approval and does not satisfy the human review requirement. Findings may be incomplete or wrong; push back on anything that looks off.

Summary

This head adds the two things the previous automated pass asked for, and both landed cleanly. normalize_after_attention now clears current_result_ended_with_done_in_proc, closing the drift against the finish_command reset block. And ATTENTION settlement no longer runs the general drain side-effect path: apply_attention_side_effect applies only ENVCHANGE, SESSIONSTATE, INFO and fatal-error liveness, so a fallible finalize_return_value can no longer retire a stream that just proved itself by reaching DONE_ATTN. The new interrupted_read_settled one-shot marker then carries that proof up to retire_after_failed_drain, which is what fixes the Copilot bot's finding that drain_stream_or_retire and close_query still discarded a synchronized session. One suggestion below, about a token the narrowed replay now drops.

Verification

Reviewed the full dd96d8e..e8dcf1c6 diff, weighted toward the 178-line delta since the last automated pass (confined to tds_client.rs), plus the surrounding unchanged code: SharedStream's poll-scoped std::sync::Mutex (guards are per-poll temporaries and send_attention_packet completes its write before the preserved read is re-awaited, so there is no same-task lock reentry), read_to_attention_boundary, push_return_value, record_error_token, observe_response_token, and every retire_after_failed_drain call site.

I traced the interrupted_read_settled lifecycle across all seven call sites and could not construct a stale-carry path. Every route that can reach retire_after_failed_drain with a TimeoutError / OperationCancelledError first passes through an entry point that clears the flag — begin_command, close_query, send_attention_with_timeout, drain_stream, or settle_interrupted_read — and fail_rpc_terminator_look_ahead only calls it on the non-attention branch, where take_settled_drain_interruption consumes the flag and still returns false.

Ran cargo nextest run -p mssql-tds --lib -E 'test(settled) or test(attention) or test(normaliz) or test(drain)' against the assigned target dir: 77 tests, all pass. I then mutated the mechanism the new tests are named for — dropping the early return from take_settled_drain_interruption in retire_after_failed_drain — and both drain_stream_or_retire_preserves_settled_cancellation and close_query_preserves_settled_cancellation fail on assertion failed: !client.transport.connection_known_dead(), while the rest stay green. They pin their own mechanism. I did not run any live-SQL suite; the ADO macOS build/test jobs and coverage-report were still pending on e8dcf1c6 when I looked, so the unchecked cargo btest box remains unconfirmed by CI. Nothing is failing.

Findings

Blocking

None.

Suggestion

mssql-tds/src/connection/tds_client.rs:4170 — narrowing the replay also drops the sp_prepexec handle.

apply_attention_side_effect has no Tokens::ReturnValue arm, so an ATTENTION drain that consumes a RETURNVALUE discards it. The previous apply_drain_side_effect routed it through finalize_return_value → push_return_value, and push_return_value is the one place that diverts RETURNVALUE ordinal 0 into prepared_handles while pending_capture is armed. normalize_after_attention clears return_values but not prepared_handles, so that capture used to survive.

The reachable case is cancelling or timing out an sp_prepexec mid-rows: SQL Server emits output parameters after the result sets, so the @handle RETURNVALUE lands inside the drain. With this change the client never learns the handle, abort_pending_prepare_capture() forgets it was preparing, and the server-side prepared handle is never sp_unprepared — it leaks for the life of the connection. Not a correctness break: the id is absent from the map, so a later execute_prepared re-prepares.

This is not a regression — main's drain_to_attention_ack discarded every token too, so it matches pre-PR behavior, and I would not hold the PR for it. But 2e50c798 did capture it, and the concern that motivated the narrowing was specifically that finalize_return_value is fallible (it decrypts against output_param_ceks). A @handle is a plain Int with no crypto metadata, so keeping just the ordinal-0 divert — or calling push_return_value and only skipping finalize_return_value when the token carries crypto metadata — gets the fallible path out of the settlement without giving up the handle. If you'd rather leave it, a one-line note on apply_attention_side_effect saying the prepared handle is knowingly forfeited would keep the next reader from re-deriving this.

Nit

None.