fix: close signal window between decode and continuation flag set

decode() previously set _continuation_expected outside the try/except
BaseException block. An asynchronously-delivered KeyboardInterrupt
landing between decode_bytes success and the flag store consumed the
bytes without poisoning the buffer; the next decode() would silently
mis-frame the continuation frame as a top-level message.

Move the flag assignment inside the try block so any exception in
that window routes through poison(). The residual window between
flag-set and `return msg` is harmless: either the message returns
cleanly or the caller must call decode_continuation() to consume the
expected continuation frame.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: antoineleclair

src/dqlitewire/codec.py

@@ -383,6 +383,15 @@ def decode(self) -> Message | None:
         # etc., and all of them mean the stream is desynchronized.
         try:
             msg = self.decode_bytes(data)
+            # The flag store MUST be inside the try block: the bytes
+            # have been consumed, so any asynchronously-delivered
+            # exception here (KeyboardInterrupt, PyErr_SetAsyncExc)
+            # between the successful decode and the flag store would
+            # otherwise leave the stream desynchronized without
+            # poisoning — the next ``decode()`` would mis-frame the
+            # continuation frame as a top-level message (issue 233).
+            if isinstance(msg, RowsResponse) and msg.has_more:
+                self._continuation_expected = True
         except BaseException as e:
             # poison() stores Exception | None; wrap non-Exception
             # BaseException subclasses so the poison cause is still a
@@ -394,9 +403,6 @@ def decode(self) -> Message | None:
             )
             raise
 
-        if isinstance(msg, RowsResponse) and msg.has_more:
-            self._continuation_expected = True
-
         return msg
 
     def decode_bytes(self, data: bytes) -> Message:

tests/test_codec_signal_safety.py

@@ -73,6 +73,30 @@ def tracer(frame: FrameType, event: str, arg: object) -> Any:
     return tracer
 
 
+def _tracer_raising_in_after_call(
+    func_name: str,
+    callee_name: str,
+    exc: BaseException = _DEFAULT_EXC,
+) -> Any:
+    """Raise ``exc`` on the first line event in ``func_name`` AFTER a
+    call to ``callee_name`` has returned. Used to target the window
+    between a sub-call and the post-call statements in the caller.
+    """
+    state = {"saw_return": False, "raised": False}
+
+    def tracer(frame: FrameType, event: str, arg: object) -> Any:
+        if state["raised"]:
+            return tracer
+        if event == "return" and frame.f_code.co_name == callee_name:
+            state["saw_return"] = True
+        elif event == "line" and frame.f_code.co_name == func_name and state["saw_return"]:
+            state["raised"] = True
+            raise exc
+        return tracer
+
+    return tracer
+
+
 class TestDecodeSignalSafety:
     """Regression tests for issue 045 (consumed-but-unpoisoned window
     in ``decode()`` / ``decode_continuation()``).
@@ -122,6 +146,42 @@ def test_decode_does_not_poison_on_oversized_header(self) -> None:
 
         assert not dec.is_poisoned
 
+    def test_keyboard_interrupt_between_parse_and_flag_set_poisons(self) -> None:
+        """233: a KeyboardInterrupt delivered in ``decode()`` AFTER
+        ``decode_bytes`` returns but BEFORE ``_continuation_expected``
+        is set used to propagate without poisoning, because the flag
+        assignment lived outside the ``except BaseException`` block.
+        The next ``decode()`` call would then read the continuation
+        frame as a top-level message (silent stream desync).
+
+        The fix moves the flag store inside the try/except block, so
+        any exception in that window poisons the buffer.
+        """
+        # Feed any valid message — a DbResponse is fine because the
+        # flag-check line runs before the isinstance() short-circuits.
+        dec = MessageDecoder(is_request=False)
+        dec.feed(_make_db(1) + _make_db(2))
+
+        tracer = _tracer_raising_in_after_call("decode", "decode_bytes")
+
+        sys.settrace(tracer)
+        try:
+            with contextlib.suppress(KeyboardInterrupt):
+                dec.decode()
+        finally:
+            sys.settrace(None)
+
+        assert dec.is_poisoned, (
+            "decode() must poison when an interrupt lands after "
+            "decode_bytes returns but before the flag store"
+        )
+
+        # The retry must fail fast with ProtocolError rather than
+        # silently returning the next message from a desynchronized
+        # stream.
+        with pytest.raises(ProtocolError, match="poisoned"):
+            dec.decode()
+
     def test_keyboard_interrupt_in_decode_continuation_poisons(self) -> None:
         """Same hazard as decode(): decode_continuation() also needs
         BaseException handling.