Let the store connection release the streams its scope outlives

src/agent_comms/store_files.py

@@ -11,6 +11,7 @@
 import tempfile
 import threading
 import time
+import weakref
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Callable, Iterable, Iterator
 from contextlib import (
@@ -447,6 +448,54 @@ class DependentStateTier(StoreLockTier):
     __slots__ = ()
 
 
+class LentCursor(sqlite3.Cursor):
+    """A cursor its connection lent out and retires before that connection closes."""
+
+    def retire(self) -> None:
+        self.close()
+        self.__class__ = RetiredCursor
+
+
+class RetiredCursor(LentCursor):
+    """Released by its connection's close; a stream finishing later has nothing to release."""
+
+    def close(self) -> None:
+        pass
+
+
+class LendingConnection(sqlite3.Connection):
+    """The connection owns every cursor it lends, so closing it releases them first.
+
+    A typed stream (TypedRow.iterate) holds its cursor until the stream is closed; a
+    caller that drops a half-read stream leaves that to the collector, after the store's
+    scope already closed this connection. Retiring the cursors here makes the scope the
+    one owner of query lifetime instead of every caller wrapping its stream in closing().
+    """
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._lent: weakref.WeakSet[LentCursor] = weakref.WeakSet()
+
+    def cursor(self, factory: type[LentCursor] = LentCursor) -> LentCursor:
+        lent = super().cursor(factory)
+        self._lent.add(lent)
+        return lent
+
+    def execute(self, sql: str, parameters=(), /) -> LentCursor:
+        return self.cursor().execute(sql, parameters)
+
+    def executemany(self, sql: str, parameters, /) -> LentCursor:
+        return self.cursor().executemany(sql, parameters)
+
+    def executescript(self, script: str, /) -> LentCursor:
+        return self.cursor().executescript(script)
+
+    def close(self) -> None:
+        for lent in tuple(self._lent):
+            lent.retire()
+        super().close()
+
+
 class StoreContention(ABC):
     """Typed SQLite busy behavior shared by snapshots and writers."""
 
@@ -455,7 +504,9 @@ class StoreContention(ABC):
     def timeout(self) -> float: ...
 
     def reader(self, uri: str) -> sqlite3.Connection:
-        return sqlite3.connect(uri, uri=True, isolation_level=None, timeout=self.timeout)
+        return sqlite3.connect(
+            uri, uri=True, isolation_level=None, timeout=self.timeout,
+            factory=LendingConnection)
 
     def writer(
         self, path: str | os.PathLike[str], *,
@@ -464,7 +515,7 @@ def writer(
     ) -> sqlite3.Connection:
         return sqlite3.connect(
             path, timeout=self.timeout, isolation_level=isolation_level,
-            uri=uri, check_same_thread=check_same_thread,
+            uri=uri, check_same_thread=check_same_thread, factory=LendingConnection,
         )
 
 

tests/round2/test_typed_table.py

@@ -1,8 +1,10 @@
 """Declaration families round-trip through real SQLite without a schema roster."""
 
+import gc
 import json
 import sqlite3
-from contextlib import nullcontext
+import sys
+from contextlib import closing, nullcontext
 from dataclasses import dataclass, field, replace
 from enum import Enum, IntFlag
 from time import perf_counter
@@ -13,6 +15,7 @@
 from agent_comms.field_codec import (
     AttemptId, BoundedText, FieldCodec, GitCommitText, GoalId, Id128Text, MalformedRecord, NoAttemptId, NullMember,
 )
+from agent_comms.store_files import DEFAULT_STORE_CONTENTION
 from agent_comms.messages import InfoKind, MessageKind
 from agent_comms.typed_table import (
     Column,
@@ -568,3 +571,18 @@ class StreamProjection(TypedRow):
         rows = StreamProjection.iterate(db.execute("SELECT 2 AS enabled"))
         with pytest.raises(ValueError):
             next(rows)
+
+
+def test_store_connection_releases_streams_its_scope_outlives(monkeypatch):
+    @dataclass(frozen=True)
+    class StreamProjection(TypedRow):
+        enabled: bool
+
+    unraisable = []
+    monkeypatch.setattr(sys, "unraisablehook", unraisable.append)
+    with closing(DEFAULT_STORE_CONTENTION.writer(":memory:")) as db:
+        rows = StreamProjection.iterate(db.execute("SELECT 1 AS enabled UNION ALL SELECT 2"))
+        assert next(rows) == StreamProjection(True)
+    del rows
+    gc.collect()
+    assert unraisable == []