fix(memory2): drain callbacks before store close

[omarespejel]

Contribution path

• Small, safe change that does not need a tracking issue
• Linked issue: Data collection x mem2 fixes #2572 — shutdown race

Problem

Memory2 could close SQLite while an admitted input or TF callback was still resolving a timestamp or pose, or appending data. Recorder setup and rollback could also cross shutdown before cleanup ownership was established, and teardown could wait indefinitely when a callback needed the store lock.

Solution

Keep the store separate from generic module disposables and close it last. Shutdown blocks callback admission, unsubscribes ordinary inputs and TF, drains admitted callbacks, disposes dispatchers, completes generic teardown, and only then closes SQLite.

Input and TF drains share one two-second deadline. If any admitted callback does not finish, shutdown raises, marks teardown failed, and leaves SQLite open. Cleanup continues after other errors; a drain failure takes precedence while retaining the earlier error as context.

The deadline fits inside DimOS's five-second process shutdown grace. Setup rollback follows the same ordering and is serialized with normal stop. If dispatcher setup finishes after shutdown wins, its disposable is cleaned up and setup is rejected.

Normal stop still drains while _memory_stop_lock is held. A callback that needs that lock therefore takes the bounded fail-closed path. Moving the drain outside the lock is a separate lifecycle change.

TeleopRecorder uses the same protected store lifecycle. Reception timestamps, poseless streams, remapping, recording contents, and robot-control behavior are unchanged.

How to Test

uv run pytest -q dimos/memory2/test_module.py dimos/memory2/test_recorder_setup_rollback.py
uv run pytest -q dimos/memory2 --ignore=dimos/memory2/codecs/test_codecs.py
uv run pytest -q dimos/teleop

Validated locally at 6285b405c:

• focused lifecycle tests: 44 passed
• Memory2 suite: 352 passed, 126 skipped
• teleop suite: 102 passed
• Ruff format/check, changed-file mypy, pre-commit hooks, commit-signature filtering, and git diff --check: passed

The private-LFS codec tests were excluded because their credentials are unavailable and they are unrelated to this change.

AI assistance

Used GPT-5 for implementation and tests, and Qodo, CodeRabbit, and Greptile for review.

Checklist

• This PR is scoped to one issue or clearly stated problem.
• I ran the relevant checks for the files I changed.
• I have reviewed and understood every line in this PR.
• I disclosed AI assistance above.
• I have read and approved the CLA.

[greptile-apps]

Greptile Summary

This PR makes Memory2 shutdown wait for recorder work before closing SQLite.

• Adds guarded store creation and shutdown state.
• Drains active input and TF callbacks during recorder teardown.
• Handles late subscription setup during shutdown.
• Uses the protected store lifecycle in TeleopRecorder.
• Adds lifecycle, error-path, and serialization tests.

Confidence Score: 5/5

This looks safe to merge.

• No blocking issues found in the changed code.

Important Files Changed

Filename	Overview
dimos/memory2/module.py	Adds coordinated store shutdown and callback draining for recorder inputs and TF subscriptions.
dimos/memory2/test_module.py	Adds deterministic coverage for shutdown ordering, callback admission, setup handoffs, and cleanup errors.
dimos/teleop/utils/recorder.py	Opens TeleopRecorder's timestamped database through the shared store lifecycle.
dimos/teleop/utils/test_recorder.py	Updates TeleopRecorder coverage for the shared store lifecycle.

_{Reviews (8): Last reviewed commit: "fix(memory2): keep store open after tear..." | Re-trigger Greptile}

[omarespejel]

I love you @greptile-apps ❤️

[greptile-apps]

Thanks! Great work on this PR — the drain-before-close approach is solid and the fix is well-scoped. 🤝

[omarespejel]

Thanks! @greptile-apps why is this pr important?

[greptile-apps]

This PR fixes a real data corruption / crash bug in the Memory2 shutdown path, which matters for a few reasons:

1. Silent data corruption risk: Before this fix, an in-flight TF (transform) callback could resume after the SQLite store was already closed and call tf_stream.append() on a closed store. Depending on the SQLite driver, this can silently corrupt the database, raise an unhandled exception in a background thread, or drop data — all bad outcomes for a system that stores robot sensor/transform history.

2. Race is non-deterministic: The bug only manifests under timing pressure (e.g., high callback frequency at shutdown, slow machines, or heavy load). This makes it the worst kind of bug — it passes tests most of the time and fails in production under stress.

3. TF setup / shutdown overlap: A second race existed where setup_tf_recording() could begin constructing a new transport while shutdown was already in progress, potentially leaving the system in a partially-initialized state with no way to clean it up.

4. Affects TeleopRecorder too: The same store lifecycle is shared by TeleopRecorder, so the fix covers teleop recording teardown as well — relevant any time a teleoperation session ends.

In short: without this fix, graceful shutdown of a robot session has undefined behavior around its persistent storage layer. That's a correctness guarantee you need before relying on Memory2 data for replay, debugging, or training.

[omarespejel]

Hey hey @leshy @paul-nechifor @spomichter, could someone approve the pending workflows and review? Both are mergeable and up to date with main. thanks!

[omarespejel]

Fixed. Shutdown now blocks new recorder callbacks, unsubscribes, lets anything already resolving a pose or writing to SQLite finish, then cancels the dispatcher and closes the store. I added tests for in-flight work and setup finishing during shutdown. @greptile-apps mind taking another look?

[omarespejel]

Updated the remaining shutdown cases. Recorder inputs and TF now drain against one shared two-second deadline. If a callback cannot finish, teardown fails closed and leaves SQLite open. This also covers setup rollback and dispatcher setup racing with stop. Added deterministic regression coverage for those paths.

@greptile-apps could you take another look at the latest commit?