fix(sqlite): isolate thread-local connections by database#472
Closed
dplush wants to merge 1 commit into
Closed
Conversation
📝 WalkthroughWalkthroughThe PR replaces thread-local singleton SQLite connections with canonical per-database connection caches and ownership tracking. Mnemosyne and BeamMemory now acquire and release shared connections explicitly, checkpoint and close them after the final owner exits, and test cleanup covers all cached database paths. ChangesSQLite connection lifecycle
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Mnemosyne
participant BeamMemory
participant ConnectionCache
participant SQLite
Mnemosyne->>ConnectionCache: Acquire connection for db_path
ConnectionCache->>SQLite: Create or validate cached connection
Mnemosyne->>BeamMemory: Initialize managed BeamMemory
BeamMemory->>ConnectionCache: Acquire BeamMemory ownership
Mnemosyne->>ConnectionCache: Release wrapper ownership on close
BeamMemory->>ConnectionCache: Release BeamMemory ownership
ConnectionCache->>SQLite: Checkpoint and close after final release
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@mnemosyne/core/beam.py`:
- Around line 469-478: In the cached connection validation in
mnemosyne/core/beam.py lines 469-478, recover only when the probe confirms the
handle is closed; propagate transient SQLite errors instead of evicting the
cache, and explicitly close the invalid connection before removal. Apply the
same narrow recovery and cleanup to the corresponding cache logic in
mnemosyne/core/memory.py lines 85-90, preserving a single connection generation
for arbitrary probe failures.
- Around line 508-526: Replace path-only ownership tracking with exact lease
tokens containing the connection generation and acquisition thread. In
mnemosyne/core/beam.py lines 508-526, have acquisition return and release
validate the stored BEAM lease; update BeamMemory.close around lines 3094-3101
to release that lease instead of self.db_path. Apply the same generation- and
thread-aware lease tracking in mnemosyne/core/memory.py lines 110-128, and make
the wrapper release its stored core lease around lines 291-293.
In `@mnemosyne/core/memory.py`:
- Around line 258-280: Update BeamMemory initialization around
_acquire_connection and BeamMemory construction so BEAM is initialized through
its ownership-aware wrapper lease before any failure-prone setup such as init_db
or get_config. Ensure the exception path releases both core and BEAM resources
when acquired, preventing a cached BEAM handle from remaining with zero owners;
preserve existing cleanup for the core connection.
In `@tests/test_beam.py`:
- Around line 955-971: Expand
test_close_isolated_by_database_path_and_live_owner and related lifecycle tests
to cover standalone BeamMemory owners on same and different database paths,
asserting both core and BEAM handles are closed after the final owner exits. Add
coverage for cross-thread close behavior according to the supported contract,
force-close followed by reopen and stale-wrapper close, and constructor failures
verifying both module caches have no remaining owners.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: d80f6c93-4593-49ba-a8e7-cae004a6da80
📒 Files selected for processing (4)
mnemosyne/core/beam.pymnemosyne/core/memory.pytests/conftest.pytests/test_beam.py
Comment on lines
+469
to
+478
| key, path = _canonical_db_path(db_path) | ||
| connections = _connections() | ||
| conn = connections.get(key) | ||
| if conn is not None: | ||
| # Verify the cached connection is still alive | ||
| try: | ||
| _thread_local.conn.execute("SELECT 1") | ||
| conn.execute("SELECT 1") | ||
| except Exception: | ||
| needs_reconnect = True | ||
| connections.pop(key, None) | ||
| conn = None |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Arbitrary probe failures must not create a second connection generation.
Both caches evict without closing on any SELECT 1 exception, potentially leaking the original handle and splitting existing wrappers from future acquisitions.
mnemosyne/core/beam.py#L469-L478: recover only from a confirmed closed handle; propagate transient SQLite errors.mnemosyne/core/memory.py#L85-L90: apply the same narrow recovery and explicit cleanup.
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 476-476: Do not catch blind exception: Exception
(BLE001)
📍 Affects 2 files
mnemosyne/core/beam.py#L469-L478(this comment)mnemosyne/core/memory.py#L85-L90
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mnemosyne/core/beam.py` around lines 469 - 478, In the cached connection
validation in mnemosyne/core/beam.py lines 469-478, recover only when the probe
confirms the handle is closed; propagate transient SQLite errors instead of
evicting the cache, and explicitly close the invalid connection before removal.
Apply the same narrow recovery and cleanup to the corresponding cache logic in
mnemosyne/core/memory.py lines 85-90, preserving a single connection generation
for arbitrary probe failures.
Source: Linters/SAST tools
Comment on lines
+508
to
+526
| def _acquire_beam_connection(db_path: Path = None) -> sqlite3.Connection: | ||
| """Get a connection and register one live BeamMemory wrapper for its path.""" | ||
| key, _ = _canonical_db_path(db_path) | ||
| conn = _get_connection(db_path) | ||
| owners = _owners() | ||
| owners[key] = owners.get(key, 0) + 1 | ||
| return conn | ||
|
|
||
|
|
||
| def _release_beam_connection(db_path: Path = None) -> None: | ||
| """Release one wrapper; close its connection only after its final owner.""" | ||
| key, path = _canonical_db_path(db_path) | ||
| owners = _owners() | ||
| remaining = owners.get(key, 0) - 1 | ||
| if remaining > 0: | ||
| owners[key] = remaining | ||
| return | ||
| owners.pop(key, None) | ||
| _close_beam_connection(path) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Path-only ownership can close another live wrapper’s connection.
The release operation must identify the exact connection generation and acquisition thread, rather than consulting whichever thread-local path entry exists when close() runs.
mnemosyne/core/beam.py#L508-L526: return/store an exact BEAM lease token and validate it during release.mnemosyne/core/beam.py#L3094-L3101: release the stored lease instead of onlyself.db_path.mnemosyne/core/memory.py#L110-L128: make core ownership generation- and thread-aware.mnemosyne/core/memory.py#L291-L293: release the wrapper’s stored core lease.
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 508-508: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
[warning] 517-517: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
📍 Affects 2 files
mnemosyne/core/beam.py#L508-L526(this comment)mnemosyne/core/beam.py#L3094-L3101mnemosyne/core/memory.py#L110-L128mnemosyne/core/memory.py#L291-L293
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mnemosyne/core/beam.py` around lines 508 - 526, Replace path-only ownership
tracking with exact lease tokens containing the connection generation and
acquisition thread. In mnemosyne/core/beam.py lines 508-526, have acquisition
return and release validate the stored BEAM lease; update BeamMemory.close
around lines 3094-3101 to release that lease instead of self.db_path. Apply the
same generation- and thread-aware lease tracking in mnemosyne/core/memory.py
lines 110-128, and make the wrapper release its stored core lease around lines
291-293.
Comment on lines
+258
to
+280
| try: | ||
| self.conn = _acquire_connection(self.db_path) | ||
| self._owns_connection = True | ||
| init_db(self.db_path) | ||
|
|
||
| # Phase 8: Streaming + Patterns + Plugins (lazy init) | ||
| self._stream = None | ||
| self._compressor = None | ||
| self._pattern_detector = None | ||
| self._delta_sync = None | ||
| self._plugin_manager = None | ||
|
|
||
| # Create beam with streaming emitter wired | ||
| self.beam = BeamMemory(session_id=session_id, db_path=self.db_path, | ||
| author_id=author_id, author_type=author_type, | ||
| channel_id=channel_id, | ||
| event_emitter=self._stream_emit) | ||
| except BaseException: | ||
| if self._owns_connection: | ||
| self._owns_connection = False | ||
| _release_connection(self.db_path) | ||
| self._closed = True | ||
| raise |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Avoid creating an unowned BEAM connection before guarded acquisition.
init_db() calls init_beam(), which can populate BEAM’s cache before BeamMemory acquires ownership. If BeamMemory then fails before _acquire_beam_connection()—for example in get_config()—this handler releases only the core connection, leaving the BEAM handle cached with zero owners.
Initialize BEAM under its wrapper lease, or add ownership-aware rollback for this path.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mnemosyne/core/memory.py` around lines 258 - 280, Update BeamMemory
initialization around _acquire_connection and BeamMemory construction so BEAM is
initialized through its ownership-aware wrapper lease before any failure-prone
setup such as init_db or get_config. Ensure the exception path releases both
core and BEAM resources when acquired, preventing a cached BEAM handle from
remaining with zero owners; preserve existing cleanup for the core connection.
Comment on lines
+955
to
+971
| def test_close_isolated_by_database_path_and_live_owner(self, temp_db, tmp_path): | ||
| other_db = tmp_path / "other.db" | ||
| first = Mnemosyne(session_id="first", db_path=temp_db) | ||
| second = Mnemosyne(session_id="second", db_path=temp_db) | ||
| other = Mnemosyne(session_id="other", db_path=other_db) | ||
|
|
||
| first.close() | ||
| assert second.conn.execute("SELECT 1").fetchone()[0] == 1 | ||
| assert second.beam.conn.execute("SELECT 1").fetchone()[0] == 1 | ||
| assert other.conn.execute("SELECT 1").fetchone()[0] == 1 | ||
| assert other.beam.conn.execute("SELECT 1").fetchone()[0] == 1 | ||
|
|
||
| second.close() | ||
| with pytest.raises(sqlite3.ProgrammingError): | ||
| second.conn.execute("SELECT 1") | ||
| assert other.conn.execute("SELECT 1").fetchone()[0] == 1 | ||
| other.close() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Expand the ownership lifecycle matrix before merge.
Add assertions for:
- Standalone
BeamMemorysame-path/different-path owners. - Both core and BEAM handles after the final owner closes.
- Cross-thread close, either safely supported or explicitly rejected.
- Force-close → reopen → stale-wrapper close.
- Constructor failure leaving both module caches owner-free.
These scenarios directly exercise the new lease boundaries rather than only the basic same-thread path.
As per path instructions, “COMPREHENSIVE REVIEW REQUIRED IN A SINGLE PASS” and “Group ALL suggestions for related fixtures/test files together.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_beam.py` around lines 955 - 971, Expand
test_close_isolated_by_database_path_and_live_owner and related lifecycle tests
to cover standalone BeamMemory owners on same and different database paths,
asserting both core and BEAM handles are closed after the final owner exits. Add
coverage for cross-thread close behavior according to the supported contract,
force-close followed by reopen and stale-wrapper close, and constructor failures
verifying both module caches have no remaining owners.
Source: Path instructions
Contributor
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Replacement for the unsafe #382 follow-up. Stores core and BEAM SQLite connections per canonical database path in thread-local maps, with per-path wrapper ownership so one live instance cannot close another database or same-database owner. Also restores the constructor initialization accidentally moved under
__del__.Verification
4 passed.py_compileandgit diff --checkpassed.Supersedes #470.
Summary
MnemosyneorBeamMemoryinstance cannot close another live owner, including same-path instances.close()and best-effort__del__().sqlite_vecsetup while improving checkpoint and test isolation.Architectural impact
This is the right call for Mnemosyne’s local-first core: it strengthens local SQLite reliability without introducing remote dependencies or weakening privacy guarantees. The change is limited to connection lifecycle management and does not alter working, episodic, or BEAM memory semantics, retrieval strategies, consolidation, veracity, synchronization, or benchmark methodology.
Hermes, MCP, and CLI integration surfaces remain behaviorally unchanged; they benefit indirectly from safer concurrent wrapper usage and deterministic cleanup. The per-path ownership model improves long-term maintainability by making resource ownership explicit and preventing cross-instance shutdown bugs.
Validation covered same-path/different-path smoke tests, a focused four-test integration suite, compilation, and whitespace checks.