fix: enforce ROWS continuation drain invariant via state flag

Add a `_continuation_expected` flag to `MessageDecoder`. When
`decode()` returns a `RowsResponse` with `has_more=True`, the flag
is set. While set, subsequent `decode()` calls raise `ProtocolError`
— the next frame in the buffer is a ROWS continuation (no column
header prefix), so `decode()` would misparse it. Callers must use
`decode_continuation()` until `has_more` is `False`, or `reset()`.

This matches Go's `protocol.Protocol` pattern where `Call()` holds
the mutex across the entire ROWS response sequence, and `Rows.Close()`
drains remaining frames before releasing the connection. The flag
converts silent stream desynchronization (decode reads continuation
data as a new message) into a loud `ProtocolError` at the point of
misuse.

The flag is cleared by `decode_continuation()` when the final frame
(`has_more=False`) is decoded, and by `reset()` for abandoning the
stream.

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

Author: antoineleclair

src/dqlitewire/codec.py

@@ -180,6 +180,7 @@ def __init__(self, is_request: bool = False, version: int = PROTOCOL_VERSION) ->
         # handshake before decoding any messages. Response decoders (client-side)
         # don't receive an inbound handshake, so they skip this check.
         self._handshake_done = not is_request
+        self._continuation_expected = False
 
     @property
     def version(self) -> int | None:
@@ -204,6 +205,7 @@ def reset(self) -> None:
         handshake state to "not yet received". Use this after a reconnect.
         """
         self._buffer.reset()
+        self._continuation_expected = False
         if self._is_request:
             self._handshake_done = False
             self._version = None
@@ -280,9 +282,12 @@ def decode_continuation(
                     f"got type {header.msg_type}"
                 )
 
-            return RowsResponse.decode_rows_continuation(
+            result = RowsResponse.decode_rows_continuation(
                 body, column_names, column_count, max_rows=max_rows
             )
+            if not result.has_more:
+                self._continuation_expected = False
+            return result
         except BaseException as e:
             # poison() stores Exception | None; wrap non-Exception
             # BaseException subclasses so the poison cause is still a
@@ -300,8 +305,17 @@ def decode(self) -> Message | None:
         Returns None if no complete message is available.
         Raises ProtocolError if called on a request decoder before decode_handshake().
         Raises ProtocolError if the decoder is poisoned.
+        Raises ProtocolError if a ROWS continuation is in progress
+        (call ``decode_continuation()`` until ``has_more`` is ``False``,
+        or ``reset()`` to abandon the stream).
         """
         self._buffer._check_poisoned()
+        if self._continuation_expected:
+            raise ProtocolError(
+                "Cannot decode a new message while a ROWS continuation "
+                "is in progress. Call decode_continuation() until "
+                "has_more is False, or call reset() to abandon the stream."
+            )
         if not self._handshake_done:
             raise ProtocolError(
                 "Protocol handshake not yet received. Call decode_handshake() before decode()."
@@ -321,7 +335,7 @@ def decode(self) -> Message | None:
         # struct.error, ValueError, UnicodeDecodeError, IndexError,
         # etc., and all of them mean the stream is desynchronized.
         try:
-            return self.decode_bytes(data)
+            msg = self.decode_bytes(data)
         except BaseException as e:
             # poison() stores Exception | None; wrap non-Exception
             # BaseException subclasses so the poison cause is still a
@@ -333,6 +347,11 @@ 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:
         """Decode a message from bytes.
 

tests/test_codec.py

@@ -939,6 +939,149 @@ def test_decode_continuation_raises_on_unexpected_type(self) -> None:
             decoder.decode_continuation(column_names=["id"], column_count=1)
 
 
+class TestDecoderContinuationExpected:
+    """Regression tests for issue 058.
+
+    When ``decode()`` returns a ``RowsResponse`` with ``has_more=True``,
+    the decoder enters a "continuation expected" state. Calling ``decode()``
+    again before draining all continuations via ``decode_continuation()``
+    is a protocol error — the next frame in the buffer is a continuation
+    (no column header prefix), so ``decode()`` would misparse it. The
+    ``_continuation_expected`` flag makes this misuse fail loudly with
+    ``ProtocolError`` instead of producing silent stream desynchronization.
+    """
+
+    def test_decode_raises_when_continuation_expected(self) -> None:
+        """decode() must refuse while a ROWS continuation is in progress."""
+        from dqlitewire.constants import ROW_PART_MARKER, ValueType
+        from dqlitewire.exceptions import ProtocolError
+        from dqlitewire.messages.base import Header
+        from dqlitewire.messages.responses import RowsResponse
+        from dqlitewire.tuples import encode_row_header, encode_row_values
+        from dqlitewire.types import encode_text, encode_uint64
+
+        types = [ValueType.INTEGER]
+
+        # Build a RowsResponse with has_more=True
+        body = encode_uint64(1)  # column_count
+        body += encode_text("id")
+        body += encode_row_header(types)
+        body += encode_row_values([1], types)
+        body += encode_uint64(ROW_PART_MARKER)
+        header = Header(size_words=len(body) // 8, msg_type=7, schema=0)
+        msg_bytes = header.encode() + body
+
+        # Also feed a second (standalone) message that decode() would try to read
+        second = ResultResponse(last_insert_id=0, rows_affected=0).encode()
+
+        decoder = MessageDecoder(is_request=False)
+        decoder.feed(msg_bytes + second)
+
+        # First decode() returns the initial RowsResponse with has_more=True
+        result = decoder.decode()
+        assert isinstance(result, RowsResponse)
+        assert result.has_more is True
+
+        # Second decode() must raise — we're in "continuation expected" state
+        with pytest.raises(ProtocolError, match="continuation"):
+            decoder.decode()
+
+    def test_decode_continuation_clears_flag(self) -> None:
+        """After draining all continuations (has_more=False), decode() works again."""
+        from dqlitewire.constants import ROW_DONE_MARKER, ROW_PART_MARKER, ValueType
+        from dqlitewire.messages.base import Header
+        from dqlitewire.messages.responses import RowsResponse
+        from dqlitewire.tuples import encode_row_header, encode_row_values
+        from dqlitewire.types import encode_text, encode_uint64
+
+        types = [ValueType.INTEGER]
+
+        # Initial frame (has_more=True)
+        body1 = encode_uint64(1)
+        body1 += encode_text("id")
+        body1 += encode_row_header(types)
+        body1 += encode_row_values([1], types)
+        body1 += encode_uint64(ROW_PART_MARKER)
+        h1 = Header(size_words=len(body1) // 8, msg_type=7, schema=0)
+
+        # Continuation frame (has_more=False)
+        body2 = encode_row_header(types)
+        body2 += encode_row_values([2], types)
+        body2 += encode_uint64(ROW_DONE_MARKER)
+        h2 = Header(size_words=len(body2) // 8, msg_type=7, schema=0)
+
+        # Normal message after the ROWS sequence
+        normal = ResultResponse(last_insert_id=5, rows_affected=3).encode()
+
+        decoder = MessageDecoder(is_request=False)
+        decoder.feed(h1.encode() + body1 + h2.encode() + body2 + normal)
+
+        # decode initial
+        initial = decoder.decode()
+        assert isinstance(initial, RowsResponse) and initial.has_more
+
+        # decode continuation
+        cont = decoder.decode_continuation(
+            column_names=initial.column_names,
+            column_count=len(initial.column_names),
+        )
+        assert isinstance(cont, RowsResponse) and not cont.has_more
+
+        # Now decode() should work again
+        result = decoder.decode()
+        assert isinstance(result, ResultResponse)
+        assert result.last_insert_id == 5
+
+    def test_reset_clears_continuation_expected(self) -> None:
+        """reset() must clear the continuation-expected flag."""
+        from dqlitewire.constants import ROW_PART_MARKER, ValueType
+        from dqlitewire.messages.base import Header
+        from dqlitewire.messages.responses import RowsResponse
+        from dqlitewire.tuples import encode_row_header, encode_row_values
+        from dqlitewire.types import encode_text, encode_uint64
+
+        types = [ValueType.INTEGER]
+        body = encode_uint64(1)
+        body += encode_text("id")
+        body += encode_row_header(types)
+        body += encode_row_values([1], types)
+        body += encode_uint64(ROW_PART_MARKER)
+        header = Header(size_words=len(body) // 8, msg_type=7, schema=0)
+
+        decoder = MessageDecoder(is_request=False)
+        decoder.feed(header.encode() + body)
+        result = decoder.decode()
+        assert isinstance(result, RowsResponse) and result.has_more
+
+        # Reset should clear the flag
+        decoder.reset()
+        normal = ResultResponse(last_insert_id=0, rows_affected=0).encode()
+        decoder.feed(normal)
+        msg = decoder.decode()
+        assert isinstance(msg, ResultResponse)
+
+    def test_has_more_false_does_not_set_flag(self) -> None:
+        """A RowsResponse with has_more=False should NOT set the flag."""
+        from dqlitewire.messages.responses import RowsResponse
+
+        decoder = MessageDecoder(is_request=False)
+        msg = RowsResponse(
+            column_names=["x"],
+            column_types=[1],
+            rows=[[1]],
+            has_more=False,
+        )
+        normal = ResultResponse(last_insert_id=0, rows_affected=0).encode()
+        decoder.feed(msg.encode() + normal)
+
+        result = decoder.decode()
+        assert isinstance(result, RowsResponse) and not result.has_more
+
+        # decode() should work fine — no continuation expected
+        result2 = decoder.decode()
+        assert isinstance(result2, ResultResponse)
+
+
 class TestDecoderSkipMessage:
     """Test skip_message() and is_skipping on MessageDecoder."""