fix: received data not fully acknowledged during h2 stream cleanup

[quettabit]

No description provided.

[greptile-apps]

Greptile Summary

This PR fixes a race condition in H2 stream cleanup where received data could fail to be acknowledged. The fix moves the unacked_flow_bytes read-and-zero operation inside _write_lock by introducing a new Connection.ack_all_data method, replacing the standalone _take_all_unacked_flow_bytes helper and the separate ack_data call in both unary_request and streaming_request cleanup paths.

Confidence Score: 5/5

Safe to merge — the fix is correct and consistent across both cleanup sites with no new logic risks.

The only finding is a P2 suggestion about missing streaming-path test coverage, which is non-blocking. The core logic change is sound and the race condition fix is well-targeted.

No files require special attention.

Important Files Changed

Filename	Overview
src/s2_sdk/_client.py	Adds Connection.ack_all_data that reads/zeros state.unacked_flow_bytes under _write_lock, eliminating the race where new bytes could be received between the old synchronous read-and-zero and the lock acquisition in ack_data; both cleanup sites updated consistently.
tests/test_client.py	Updates unary timeout test to assert ack_all_data(1, state) instead of ack_data(1, 11); no corresponding test added for the streaming_request cleanup path.

Sequence Diagram

sequenceDiagram
    participant CL as Cleanup (finally)
    participant WL as _write_lock
    participant RL as RecvLoop
    participant H2 as h2 layer

    note over CL,H2: OLD behaviour (race condition)
    CL->>CL: _take_all_unacked_flow_bytes() → reads & zeroes state (no lock)
    RL-->>WL: acquires _write_lock, increments unacked_flow_bytes
    RL-->>H2: receive_data / handle events
    RL-->>WL: releases _write_lock
    CL->>WL: await ack_data() → acquires lock
    CL->>H2: acknowledge_received_data(old_nbytes) ← new bytes MISSED

    note over CL,H2: NEW behaviour (this PR)
    CL->>WL: await ack_all_data() → acquires _write_lock
    WL-->>RL: recv loop blocked until lock released
    CL->>CL: read & zero unacked_flow_bytes inside lock
    CL->>H2: acknowledge_received_data(all bytes) ← nothing missed
    CL-->>WL: releases _write_lock

Loading

Comments Outside Diff (1)

1. tests/test_client.py, line 204-232 (link)

   Missing test for streaming_request cleanup path

   Both unary_request and streaming_request were updated to call ack_all_data, but only the unary timeout path has a test. A symmetric test for streaming_request (e.g., cancellation or early-exit while data is in flight) would confirm that the streaming cleanup also correctly calls ack_all_data(stream_id, state) and reset_stream.

   Prompt To Fix With AI

   This is a comment left during a code review.
   Path: tests/test_client.py
   Line: 204-232
   
   Comment:
   **Missing test for streaming_request cleanup path**
   
   Both `unary_request` and `streaming_request` were updated to call `ack_all_data`, but only the unary timeout path has a test. A symmetric test for `streaming_request` (e.g., cancellation or early-exit while data is in flight) would confirm that the streaming cleanup also correctly calls `ack_all_data(stream_id, state)` and `reset_stream`.
   
   How can I resolve this? If you propose a fix, please make it concise.

Prompt To Fix All With AI

This is a comment left during a code review.
Path: tests/test_client.py
Line: 204-232

Comment:
**Missing test for streaming_request cleanup path**

Both `unary_request` and `streaming_request` were updated to call `ack_all_data`, but only the unary timeout path has a test. A symmetric test for `streaming_request` (e.g., cancellation or early-exit while data is in flight) would confirm that the streaming cleanup also correctly calls `ack_all_data(stream_id, state)` and `reset_stream`.

How can I resolve this? If you propose a fix, please make it concise.

_{Reviews (1): Last reviewed commit: "initial commit" | Re-trigger Greptile}