feat: key-package maintenance as TaskRunner recurring tasks

[insipx]

Stacked on #3806 (generic recurrence); the KpRotation/KpDeletion proto variants land in #3805. Replaces the 5s-poll KeyPackageCleaner worker with two recurring TaskRunner tasks.

What

• KpRotation — recurring singleton: rotates + uploads a fresh key package when the identity's rotation deadline (next_key_package_rotation_ns) is due, then reschedules to the live column (never a hardcoded +30d). On rotation it self-heals the deletion singleton and pulls it in to the earliest pending delete_at_ns.
• KpDeletion — recurring singleton: sweeps expired local key-package material, reschedules to the next pending deadline.
• Welcome nudge — after queue_key_package_rotation lowers the column (~5s debounce, a security property), a durable PullInDeadline pulls the rotation task in to match. Closes a real race: without it, a seed that dispatched before the column write parks ~30d out and the debounce is silently lost (reproduced; regression-tested).
• Builder seeds both singletons + reconciles deadlines to the live DB columns on every startup (idempotent; also repairs rows stranded by a crash between a column write and its pull-in).
• Gutted the old 5s-poll worker: loop/factory/rotate path deleted; the struct + local-deletion helpers remain as the KpDeletion arm's implementation.

Behavior notes

• KP maintenance no longer has its own toggle — it is a critical MLS function, deliberately coupled to the TaskRunner (no standalone fallback). WorkerKind::KeyPackageCleaner remains in core and all three bindings' enums for FFI compatibility, but enabling/disabling it is now a no-op. Disabling the TaskRunner itself — per-kind via worker config or globally via disable_workers — disables KP maintenance with it.
• Pre-registration (fresh client, no identity row): readers return None/false, the rotation seed parks ~30d and converges silently once registration writes the column. No transient errors.
• Zero migrations; zero reaper changes — recurring rows are never-expiring seeds (expires_at_ns = i64::MAX, max_attempts = i32::MAX), so the existing reaper predicate can't fire on them.

🤖 Generated with Claude Code

Note

Replace key-package cleaner worker with TaskRunner recurring tasks for rotation and deletion

• Removes the legacy KeyPackagesCleanerWorker and replaces it with durable KpRotation and KpDeletion task rows managed by the TaskRunner in key_package_maintenance.rs.
• KpRotation tasks reschedule to the DB-backed next_key_package_rotation_ns deadline; KpDeletion tasks reschedule to the earliest pending delete_at_ns across marked key packages.
• seed_and_reconcile_kp_tasks is called at startup (after TaskWorker registration) to seed both recurring tasks and pull in deadlines from live DB state.
• queue_key_rotation is now synchronous and atomically updates the rotation column and enqueues a PullInDeadline nudge; nudge_deletion does the same for deletion.
• is_identity_needs_rotation now returns false pre-registration and treats a NULL rotation column on an existing row as due immediately, fixing a bug where migrated DBs with NULL deadlines were skipped.
• Behavioral Change: key-package maintenance no longer runs as a standalone polling worker; it is now driven entirely by the TaskRunner dispatch loop.

Changes since #3814 opened

• Modified key rotation queueing to insert-or-ignore rotation seed tasks before enqueuing pull-in tasks [3cb7a61]
• Added tests verifying rotation seed self-healing behavior during key rotation nudges [3cb7a61]
• Changed KpDeletion task rescheduling to use rotation interval constant instead of fixed 30-day value [3cb7a61]
• Updated mock implementation to match new trait signature [3cb7a61]

^{Macroscope summarized 74a0d28.}

[claude]

Claude finished @insipx's task —— View job

---

Code Review Complete

Review Date: July 9, 2026, 12:00 AM UTC

Summary

This PR successfully refactors key-package maintenance from a polling worker to TaskRunner recurring tasks. The implementation is solid overall, with excellent test coverage and thoughtful handling of edge cases discovered during review iterations. Below are focused observations on aspects not extensively covered in prior reviews.

---

Key Observations

1. Transaction Atomicity & Recovery (✅ Well Done)

The queue_key_rotation_with_nudge implementation (identity.rs:124-185) demonstrates excellent atomicity:

• Column write + pull-in enqueue happen in one transaction
• Pre-registration (no identity row) handled gracefully as no-op via .optional()
• Target seed insertion uses insert_or_ignore to ensure pull-in never targets a missing task
• Read-back inside transaction ensures payload stability for coalescing

This addresses the partial-failure window that was previously discussed and eliminates the race condition entirely.

2. State Machine Correctness

The three-state model for next_key_package_rotation_ns is handled consistently:

• None (no identity row) → pre-registration, no rotation needed
• Some(None) (NULL column) → rotation due immediately
• Some(Some(t)) → scheduled deadline

The NULL-handling fix in queue_key_package_rotation (line 113: .or(is_null())) ensures migrated databases get the 5s debounce rather than immediate rotation.

3. Error Classification & Reconnect Logic (✅ Recently Fixed)

The IdentityError::needs_db_reconnect implementation (identity.rs:347-357) now correctly forwards connection errors from the OpenMLS keystore path:

Self::OpenMlsStorageError(SqlKeyStoreError::Connection(c)) => c.db_needs_connection()

This ensures that keystore connection pool loss during rotate/delete operations triggers worker restart rather than infinite backoff.

4. Test Coverage Strengths

The test suite in key_package_maintenance.rs demonstrates thorough coverage:

• Startup reconciliation with stale parked tasks
• NULL column initialization
• Missing seed self-healing
• Pre-registration no-op behavior
• Backoff retry convergence
• Error forwarding for DB reconnect cases

The regression tests are well-named and document the specific bugs they prevent (e.g., queue_initializes_null_rotation_column, welcome_nudge_pulls_in_parked_rotation).

---

Minor Observations (Non-Blocking)

A. Constant Usage Consistency

Lines tasks.rs:425 and :437 use KEY_PACKAGE_ROTATION_INTERVAL_NS as the park-when-idle fallback for both rotation and deletion tasks. While functionally correct (both ~30d), this reads as if the interval is load-bearing for deletion timing. Consider documenting why the same constant is appropriate for both, or extracting a KP_TASK_IDLE_PARK_INTERVAL_NS constant with a comment explaining it's just a "distant future" value that nudges will pull in sooner.

B. Builder Failure Mode

Line builder.rs:428 propagates seed_and_reconcile_kp_tasks errors with ?, making transient DB failures during startup fatal to client construction. The comment justifies this as "building a client whose critical maintenance silently never got seeded would be worse." This is a reasonable product decision but represents a behavior change from the old worker registration (which was effectively infallible). Consider adding a log line before the ? to help diagnose failures in the field:

if let Err(e) = crate::worker::key_package_maintenance::seed_and_reconcile_kp_tasks(&context) {
    tracing::error!("KP task seeding failed, cannot build client: {e}");
    return Err(e.into());
}

C. Welcome Sync Nudge Failure Handling

Line welcome_sync.rs:165-169 downgrades nudge_rotation failures to warnings. While the self-healing properties (startup reconcile, next welcome retry) mean this won't permanently strand rotation, long-lived clients with infrequent welcomes could delay rotation significantly. The comment explains this trade-off (discarding committed welcome groups would be worse), but documenting the recovery paths in the warning message would help future maintainers:

tracing::warn!(
    "key rotation nudge failed after welcome sync (will self-heal on restart or next welcome): {e}"
);

D. Documentation Opportunity

The PR description and in-code comments are excellent, but consider adding a module-level doc comment to key_package_maintenance.rs summarizing the overall design:

• Recurring tasks vs. one-shot tasks
• Nudge/pull-in mechanism
• Recovery properties (startup reconcile, retry convergence)
• Why it's coupled to TaskRunner

This would help new contributors understand the architecture without reading the entire PR thread.

---

Security & Performance

Security: No concerns. The changes reduce attack surface by eliminating the polling loop and making maintenance event-driven.

Performance: The event-driven approach is more efficient than 5s polling. The startup seeding adds negligible overhead (two task row inserts + two pull-in inserts, all insert_or_ignore).

Concurrency: The transaction-based approach in queue_key_rotation_with_nudge eliminates race conditions. The unconditional nudge_deletion call in the KpRotation arm (tasks.rs:419) ensures eventual consistency even if nudges fail transiently.

---

Conclusion

This is high-quality work with excellent attention to edge cases and recovery properties. The architectural shift from polling to event-driven maintenance is the right direction. The issues raised in earlier reviews have been addressed comprehensively. The minor observations above are suggestions for clarity and diagnostics rather than correctness concerns.

Recommendation: Approve with optional follow-up for documentation enhancements.

---

Tasks Completed

• Read and analyze changed files
• Identify novel issues not covered in prior reviews
• Provide focused, actionable feedback
• Update comment with final review

---

[codecov]

Codecov Report

❌ Patch coverage is 95.04950% with 25 lines in your changes missing coverage. Please review.
✅ Project coverage is 85.12%. Comparing base (e181466) to head (3cb7a61).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/xmtp_db/src/encrypted_store/identity.rs	90.32%	12 Missing ⚠️
...tes/xmtp_mls/src/worker/key_package_maintenance.rs	96.97%	10 Missing ⚠️
...xmtp_db/src/encrypted_store/key_package_history.rs	82.35%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3814      +/-   ##
==========================================
+ Coverage   84.74%   85.12%   +0.37%     
==========================================
  Files         409      411       +2     
  Lines       61369    63290    +1921     
==========================================
+ Hits        52008    53873    +1865     
- Misses       9361     9417      +56     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[macroscopeapp]

🟠 High src/client.rs:1078

Client::queue_key_rotation now fails with a database NotFound error when called on a client that has not yet registered. The previous implementation was a no-op in that state, but the new one delegates to queue_key_rotation_with_nudge, which reads the identity row inside the transaction and surfaces Diesel's NotFound when the row does not exist. Any caller that invokes queue_key_rotation() on a fresh client now receives an error instead of a benign no-op. Consider handling the missing-row case (e.g., treat it as a no-op or early-return) so the method remains safe to call before registration.

Also found in 1 other location(s)

crates/xmtp_mls/src/worker/key_package_maintenance.rs:177

queue_key_rotation now calls db.queue_key_rotation_with_nudge(...), which reads the identity row back with .first(...). On a client that exists before register_identity() has stored the row, that query returns NotFound, so the public Client::queue_key_rotation() API now fails instead of behaving like the old queue_key_package_rotation() path (which was an empty-table no-op).

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @crates/xmtp_mls/src/client.rs around line 1078:

`Client::queue_key_rotation` now fails with a database `NotFound` error when called on a client that has not yet registered. The previous implementation was a no-op in that state, but the new one delegates to `queue_key_rotation_with_nudge`, which reads the `identity` row inside the transaction and surfaces Diesel's `NotFound` when the row does not exist. Any caller that invokes `queue_key_rotation()` on a fresh client now receives an error instead of a benign no-op. Consider handling the missing-row case (e.g., treat it as a no-op or early-return) so the method remains safe to call before registration.

Also found in 1 other location(s):
- crates/xmtp_mls/src/worker/key_package_maintenance.rs:177 -- `queue_key_rotation` now calls `db.queue_key_rotation_with_nudge(...)`, which reads the `identity` row back with `.first(...)`. On a client that exists before `register_identity()` has stored the row, that query returns `NotFound`, so the public `Client::queue_key_rotation()` API now fails instead of behaving like the old `queue_key_package_rotation()` path (which was an empty-table no-op).

[insipx]

Good catch — real regression. Fixed in the branch tip: the in-transaction read-back now uses .optional(), and a missing identity row (pre-registration) early-returns as a no-op — no column write matched, no pull-in inserted — matching the old zero-rows-matched behavior exactly. Regression test queue_with_nudge_is_noop_before_registration pins it (empty identity table → Ok + zero task rows).

[macroscopeapp]

Approvability

Verdict: Needs human review

1 blocking correctness issue found. Architectural change replacing standalone worker with task-based key package maintenance. Unresolved high-severity comment identifies behavioral regression where pre-registration calls now error instead of being no-ops. Questions about startup error handling remain open.

^{You can customize Macroscope's approvability policy. Learn more.}

[tylerhawkes]

Replacing the KP cleaner with recurring TaskRunner tasks reads well and the failure-mode test coverage is strong. A few notes: one medium (disable coupling), two low, one nit.

[tylerhawkes]

Maintenance is gated on TaskRunner, which is individually disableable (no_runner_cfg() in tests). So disabling just the TaskRunner silently stops rotation/deletion — contradicting the "only disable_workers can kill it" claim. Guard against it, or soften the wording?

[tylerhawkes]

This ? makes a transient DB error at startup fail the whole client build; the worker registration it replaces was effectively infallible. Intended?

[tylerhawkes]

Idle deletion parks at now + NS_IN_30_DAYS; idle rotation uses KEY_PACKAGE_ROTATION_INTERVAL_NS. Both ~30d — use one constant so the gap doesn't read as load-bearing.