fix: uncaught asyncio.CancelledError in Connection.send_headers

[quettabit]

closes #81

[quettabit]

@greptileai review pls

[greptile-apps]

Greptile Summary

This PR fixes a resource leak in Connection.send_headers where a cancelled asyncio task would leave a stream ID registered in self._streams without being cleaned up. In Python 3.8+, asyncio.CancelledError is a subclass of BaseException, not Exception, so the original except Exception: guard silently missed it.

• Root cause: _flush_h2_data_and_drain() is an await-able call, meaning the task can be cancelled at that point; stream_id was registered in self._streams just before the try block but would never be cleaned up on cancellation.
• Fix: Adding asyncio.CancelledError to the except clause ensures the _streams.pop(stream_id, None) cleanup runs on cancellation and re-raises the error correctly.

Confidence Score: 5/5

Safe to merge — the change is a precise, well-scoped fix that closes a real stream-leak on task cancellation without altering any other code paths.

The fix correctly addresses the root cause: asyncio.CancelledError is a BaseException subclass in Python 3.8+ and was silently bypassing the cleanup block. The stream ID registered immediately before the try could leak permanently into _streams on cancellation. The one-line change is minimal and does not touch any other logic.

No files require special attention.

Important Files Changed

Filename	Overview
src/s2_sdk/_client.py	Single-line fix adding asyncio.CancelledError to the except clause in send_headers to ensure stream cleanup on task cancellation; logic is correct and targeted.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller
    participant send_headers
    participant _streams
    participant _flush_h2_data_and_drain

    Caller->>send_headers: await send_headers(state, headers)
    send_headers->>_streams: "_streams[stream_id] = state"
    send_headers->>_flush_h2_data_and_drain: await (can be cancelled here)
    
    alt Normal flow
        _flush_h2_data_and_drain-->>send_headers: OK
        send_headers-->>Caller: return stream_id
    else Exception raised
        _flush_h2_data_and_drain-->>send_headers: raise Exception
        send_headers->>_streams: pop(stream_id) cleanup
        send_headers-->>Caller: re-raise
    else CancelledError (BEFORE fix)
        _flush_h2_data_and_drain-->>send_headers: raise CancelledError
        Note over send_headers,_streams: except Exception misses CancelledError, stream_id leaks in _streams
        send_headers-->>Caller: propagates uncaught
    else CancelledError (AFTER fix)
        _flush_h2_data_and_drain-->>send_headers: raise CancelledError
        send_headers->>_streams: pop(stream_id) cleanup
        send_headers-->>Caller: re-raise CancelledError
    end

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Caller
    participant send_headers
    participant _streams
    participant _flush_h2_data_and_drain

    Caller->>send_headers: await send_headers(state, headers)
    send_headers->>_streams: "_streams[stream_id] = state"
    send_headers->>_flush_h2_data_and_drain: await (can be cancelled here)
    
    alt Normal flow
        _flush_h2_data_and_drain-->>send_headers: OK
        send_headers-->>Caller: return stream_id
    else Exception raised
        _flush_h2_data_and_drain-->>send_headers: raise Exception
        send_headers->>_streams: pop(stream_id) cleanup
        send_headers-->>Caller: re-raise
    else CancelledError (BEFORE fix)
        _flush_h2_data_and_drain-->>send_headers: raise CancelledError
        Note over send_headers,_streams: except Exception misses CancelledError, stream_id leaks in _streams
        send_headers-->>Caller: propagates uncaught
    else CancelledError (AFTER fix)
        _flush_h2_data_and_drain-->>send_headers: raise CancelledError
        send_headers->>_streams: pop(stream_id) cleanup
        send_headers-->>Caller: re-raise CancelledError
    end

Loading

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