From 0d3f4a9f37c437fb71f1afd9ed5f417b8e020085 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Fri, 30 Oct 2026 08:37:00 -0400 Subject: [PATCH] Let the store connection release the streams its scope outlives Every heartbeat pass logged "Exception ignored while closing generator TypedRow.iterate ... Cannot operate on a closed database": a caller dropped a half-read typed stream, the store scope closed its connection, and the collector later ran the stream's finally, closing a cursor on a dead database. StoreContention, the one place store connections are opened, now opens a LendingConnection. It owns every cursor it lends and retires them before it closes, so a stream released afterwards has nothing left to release. Callers no longer need closing() around each stream to cover for the scope. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XbFMSoWgC3ZwKzjT4SwMDu --- src/agent_comms/store_files.py | 55 ++++++++++++++++++++++++++++++-- tests/round2/test_typed_table.py | 20 +++++++++++- 2 files changed, 72 insertions(+), 3 deletions(-) diff --git a/src/agent_comms/store_files.py b/src/agent_comms/store_files.py index 36627b072b..1dd57480e9 100644 --- a/src/agent_comms/store_files.py +++ b/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, ) diff --git a/tests/round2/test_typed_table.py b/tests/round2/test_typed_table.py index e8c6dc5093..73d2264295 100644 --- a/tests/round2/test_typed_table.py +++ b/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 == []