feat: key-package maintenance as TaskRunner recurring tasks (#3814)

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](https://claude.com/claude-code)


<!-- Macroscope's pull request summary starts here -->
<!-- Macroscope will only edit the content between these invisible
markers, and the markers themselves will not be visible in the GitHub
rendered markdown. -->
<!-- If you delete either of the start / end markers from your PR's
description, Macroscope will append its summary at the bottom of the
description. -->
> [!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](https://github.com/xmtp/libxmtp/pull/3814/files#diff-1d4be0f0b8283620c52870358d20c933cc8086d19b8c804906ce1fc6d99af0eb).
> - `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.
>
> <!-- Macroscope's changelog starts here -->
> #### 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's changelog ends here -->
>
> <!-- Macroscope's review summary starts here -->
>
> <sup><a href="https://app.macroscope.com">Macroscope</a> summarized
74a0d28.</sup>
> <!-- Macroscope's review summary ends here -->
>
> <!-- macroscope-ui-refresh -->
<!-- Macroscope's pull request summary ends here -->

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: insipx

crates/xmtp_db/src/encrypted_store/identity.rs

@@ -47,11 +47,26 @@ impl StoredIdentity {
 }
 pub trait QueryIdentity {
     fn queue_key_package_rotation(&self) -> Result<(), StorageError>;
+    /// Atomically lower/initialize the rotation column (5s debounce) AND enqueue a
+    /// `PullInDeadline` task targeting `rotation_task_hash` at the resulting column
+    /// value — one transaction, so neither write can land without the other.
+    /// `rotation_seed` is insert-or-ignored first so the pull-in always has a live
+    /// target (commit-target-first), even if startup seeding never ran.
+    /// Callers wake the TaskWorker AFTER this returns (never inside a tx).
+    fn queue_key_rotation_with_nudge(
+        &self,
+        rotation_task_hash: &crate::tasks::TaskDataHash,
+        rotation_seed: crate::tasks::NewTask,
+    ) -> Result<(), StorageError>;
     fn reset_key_package_rotation_queue(
         &self,
         rotation_interval_ns: i64,
     ) -> Result<(), StorageError>;
     fn is_identity_needs_rotation(&self) -> Result<bool, StorageError>;
+    /// The identity's absolute rotation deadline (`next_key_package_rotation_ns`).
+    /// `None` if NULL or if no identity row exists yet (indistinguishable to callers;
+    /// treat as "no scheduled deadline").
+    fn next_key_package_rotation_ns(&self) -> Result<Option<i64>, StorageError>;
 }
 
 impl<T> QueryIdentity for &T
@@ -62,6 +77,14 @@ where
         (**self).queue_key_package_rotation()
     }
 
+    fn queue_key_rotation_with_nudge(
+        &self,
+        rotation_task_hash: &crate::tasks::TaskDataHash,
+        rotation_seed: crate::tasks::NewTask,
+    ) -> Result<(), StorageError> {
+        (**self).queue_key_rotation_with_nudge(rotation_task_hash, rotation_seed)
+    }
+
     fn reset_key_package_rotation_queue(
         &self,
         rotation_interval_ns: i64,
@@ -72,14 +95,24 @@ where
     fn is_identity_needs_rotation(&self) -> Result<bool, StorageError> {
         (**self).is_identity_needs_rotation()
     }
+
+    fn next_key_package_rotation_ns(&self) -> Result<Option<i64>, StorageError> {
+        (**self).next_key_package_rotation_ns()
+    }
 }
 
 impl<C: ConnectionExt> QueryIdentity for DbConnection<C> {
     fn queue_key_package_rotation(&self) -> Result<(), StorageError> {
         self.raw_query(|conn| {
             let rotate_at_ns = now_ns() + KEY_PACKAGE_QUEUE_INTERVAL_NS;
+            // NULL (migrated DBs) counts as unscheduled: initialize it here so the
+            // 5s debounce applies and nudge payloads stay stable (coalescing).
             diesel::update(dsl::identity)
-                .filter(dsl::next_key_package_rotation_ns.gt(rotate_at_ns))
+                .filter(
+                    dsl::next_key_package_rotation_ns
+                        .gt(rotate_at_ns)
+                        .or(dsl::next_key_package_rotation_ns.is_null()),
+                )
                 .set(dsl::next_key_package_rotation_ns.eq(rotate_at_ns))
                 .execute(conn)?;
 
@@ -89,6 +122,69 @@ impl<C: ConnectionExt> QueryIdentity for DbConnection<C> {
         Ok(())
     }
 
+    fn queue_key_rotation_with_nudge(
+        &self,
+        rotation_task_hash: &crate::tasks::TaskDataHash,
+        rotation_seed: crate::tasks::NewTask,
+    ) -> Result<(), StorageError> {
+        use crate::schema::tasks;
+        use diesel::Connection;
+        use xmtp_proto::xmtp::mls::database::{PullInDeadline, Task as TaskProto, task::Task};
+
+        let hash = rotation_task_hash.to_vec();
+        self.raw_query(|conn| {
+            conn.transaction::<_, diesel::result::Error, _>(|conn| {
+                let rotate_at_ns = now_ns() + KEY_PACKAGE_QUEUE_INTERVAL_NS;
+                diesel::update(dsl::identity)
+                    .filter(
+                        dsl::next_key_package_rotation_ns
+                            .gt(rotate_at_ns)
+                            .or(dsl::next_key_package_rotation_ns.is_null()),
+                    )
+                    .set(dsl::next_key_package_rotation_ns.eq(rotate_at_ns))
+                    .execute(conn)?;
+
+                // Read back inside the tx: the column is stable between rotations,
+                // so repeat calls produce byte-identical pull-ins that coalesce.
+                let deadline: Option<Option<i64>> = dsl::identity
+                    .select(dsl::next_key_package_rotation_ns)
+                    .first::<Option<i64>>(conn)
+                    .optional()?;
+                // Pre-registration (no identity row): match the old zero-rows-
+                // matched no-op instead of erroring; nothing to rotate yet.
+                let Some(deadline) = deadline else {
+                    return Ok(());
+                };
+
+                // Ensure the pull-in's target exists (no-op when already seeded):
+                // a client whose startup seeding never ran must not enqueue a
+                // dropped-on-miss nudge.
+                diesel::insert_or_ignore_into(tasks::table)
+                    .values(rotation_seed)
+                    .execute(conn)?;
+
+                let pull_in = crate::tasks::NewTask::builder()
+                    .originating_message_sequence_id(0)
+                    .originating_message_originator_id(0)
+                    .expires_at_ns(crate::tasks::NEVER_EXPIRES)
+                    .max_attempts(i32::MAX)
+                    .build(TaskProto {
+                        task: Some(Task::PullInDeadline(PullInDeadline {
+                            target_data_hash: hash,
+                            not_later_than_ns: deadline.unwrap_or(rotate_at_ns),
+                        })),
+                    })
+                    // All required builder fields are set above; unreachable.
+                    .map_err(|_| diesel::result::Error::RollbackTransaction)?;
+                diesel::insert_or_ignore_into(tasks::table)
+                    .values(pull_in)
+                    .execute(conn)?;
+                Ok(())
+            })
+        })?;
+        Ok(())
+    }
+
     fn reset_key_package_rotation_queue(
         &self,
         rotation_interval_ns: i64,
@@ -113,17 +209,33 @@ impl<C: ConnectionExt> QueryIdentity for DbConnection<C> {
     fn is_identity_needs_rotation(&self) -> Result<bool, StorageError> {
         use crate::schema::identity::dsl;
 
-        let next_rotation_opt: Option<i64> = self.raw_query(|conn| {
+        let next_rotation_opt: Option<Option<i64>> = self.raw_query(|conn| {
             dsl::identity
                 .select(dsl::next_key_package_rotation_ns)
                 .first::<Option<i64>>(conn)
+                .optional()
         })?;
 
         Ok(match next_rotation_opt {
-            Some(rotate_at) => now_ns() >= rotate_at,
-            None => true,
+            // No identity row (pre-registration): nothing to rotate yet.
+            None => false,
+            // NULL column on an existing row: rotation is due now.
+            Some(None) => true,
+            Some(Some(rotate_at)) => now_ns() >= rotate_at,
         })
     }
+
+    fn next_key_package_rotation_ns(&self) -> Result<Option<i64>, StorageError> {
+        use crate::schema::identity::dsl;
+        // Use optional() so an empty table (pre-registration) returns Ok(None).
+        let v: Option<Option<i64>> = self.raw_query(|conn| {
+            dsl::identity
+                .select(dsl::next_key_package_rotation_ns)
+                .first::<Option<i64>>(conn)
+                .optional()
+        })?;
+        Ok(v.flatten())
+    }
 }
 
 #[cfg(test)]
@@ -132,6 +244,86 @@ pub(crate) mod tests {
     use crate::{Store, XmtpTestDb};
     use xmtp_common::rand_vec;
 
+    /// A stand-in rotation seed for exercising `queue_key_rotation_with_nudge`
+    /// (the real seed payload lives in xmtp_mls).
+    fn test_rotation_seed() -> crate::tasks::NewTask {
+        use xmtp_proto::xmtp::mls::database::{KpRotation, Task as TaskProto, task::Task};
+        crate::tasks::NewTask::builder()
+            .originating_message_sequence_id(0)
+            .originating_message_originator_id(0)
+            .expires_at_ns(crate::tasks::NEVER_EXPIRES)
+            .max_attempts(i32::MAX)
+            .next_attempt_at_ns(0)
+            .build(TaskProto {
+                task: Some(Task::KpRotation(KpRotation {})),
+            })
+            .unwrap()
+    }
+
+    #[xmtp_common::test]
+    fn queue_with_nudge_is_noop_before_registration() {
+        use crate::prelude::{QueryIdentity, QueryTasks};
+        use crate::test_utils::with_connection;
+        with_connection(|conn| {
+            // Empty identity table (pre-registration): must be a no-op like the
+            // old column-only path, not a NotFound error. The seed must NOT be
+            // inserted either — pre-registration means zero writes.
+            let hash = crate::tasks::TaskDataHash::try_from([0x11u8; 32].as_slice()).unwrap();
+            conn.queue_key_rotation_with_nudge(&hash, test_rotation_seed())
+                .unwrap();
+            assert!(
+                conn.get_tasks().unwrap().is_empty(),
+                "no pull-in without an identity row"
+            );
+        })
+    }
+
+    #[xmtp_common::test]
+    fn queue_with_nudge_selfheals_missing_seed() {
+        use crate::prelude::{QueryIdentity, QueryTasks};
+        use crate::test_utils::with_connection;
+        with_connection(|conn| {
+            StoredIdentity::new("".to_string(), rand_vec::<24>(), rand_vec::<24>())
+                .store(conn)
+                .unwrap();
+            let seed = test_rotation_seed();
+            let hash = crate::tasks::TaskDataHash::try_from(seed.data_hash.as_slice()).unwrap();
+            conn.queue_key_rotation_with_nudge(&hash, seed).unwrap();
+            let tasks = conn.get_tasks().unwrap();
+            assert!(
+                tasks.iter().any(|t| t.data_hash == hash.as_ref()),
+                "nudge must insert the missing rotation seed (pull-in target)"
+            );
+            assert_eq!(tasks.len(), 2, "seed + pull-in");
+        })
+    }
+
+    #[xmtp_common::test]
+    fn queue_initializes_null_rotation_column() {
+        use crate::prelude::QueryIdentity;
+        use crate::test_utils::with_connection;
+        use xmtp_configuration::KEY_PACKAGE_QUEUE_INTERVAL_NS;
+        with_connection(|conn| {
+            StoredIdentity::new("".to_string(), rand_vec::<24>(), rand_vec::<24>())
+                .store(conn)
+                .unwrap();
+
+            // Migrated DBs have NULL here; queueing must initialize it (5s
+            // debounce) rather than skip the row.
+            conn.queue_key_package_rotation().unwrap();
+            let v = conn
+                .next_key_package_rotation_ns()
+                .unwrap()
+                .expect("NULL column must be initialized");
+            let now = xmtp_common::time::now_ns();
+            assert!(v > now && v <= now + KEY_PACKAGE_QUEUE_INTERVAL_NS);
+
+            // Lower-only: a later queue call never raises the deadline.
+            conn.queue_key_package_rotation().unwrap();
+            assert_eq!(conn.next_key_package_rotation_ns().unwrap().unwrap(), v);
+        })
+    }
+
     #[xmtp_common::test]
     async fn can_only_store_one_identity() {
         let store = crate::TestDb::create_ephemeral_store().await;

crates/xmtp_db/src/encrypted_store/key_package_history.rs

@@ -47,6 +47,10 @@ pub trait QueryKeyPackageHistory {
 
     fn get_expired_key_packages(&self) -> Result<Vec<StoredKeyPackageHistoryEntry>, StorageError>;
 
+    /// Soonest pending `delete_at_ns` across all key packages marked for deletion,
+    /// or `None` if none are marked. The KpDeletion task's reschedule source.
+    fn min_key_package_delete_at_ns(&self) -> Result<Option<i64>, StorageError>;
+
     fn delete_key_package_history_up_to_id(&self, id: i32) -> Result<(), StorageError>;
 
     fn delete_key_package_entry_with_id(&self, id: i32) -> Result<(), StorageError>;
@@ -86,6 +90,10 @@ where
         (**self).get_expired_key_packages()
     }
 
+    fn min_key_package_delete_at_ns(&self) -> Result<Option<i64>, StorageError> {
+        (**self).min_key_package_delete_at_ns()
+    }
+
     fn delete_key_package_history_up_to_id(&self, id: i32) -> Result<(), StorageError> {
         (**self).delete_key_package_history_up_to_id(id)
     }
@@ -163,6 +171,18 @@ impl<C: ConnectionExt> QueryKeyPackageHistory for DbConnection<C> {
         .map_err(StorageError::from) // convert ConnectionError into StorageError
     }
 
+    fn min_key_package_delete_at_ns(&self) -> Result<Option<i64>, StorageError> {
+        use crate::schema::key_package_history::dsl;
+        use diesel::dsl::min;
+        let v: Option<i64> = self.raw_query(|conn| {
+            dsl::key_package_history
+                .filter(dsl::delete_at_ns.is_not_null())
+                .select(min(dsl::delete_at_ns))
+                .first::<Option<i64>>(conn)
+        })?;
+        Ok(v)
+    }
+
     fn delete_key_package_history_up_to_id(&self, id: i32) -> Result<(), StorageError> {
         self.raw_query(|conn| {
             diesel::delete(
@@ -194,6 +214,14 @@ mod tests {
     use crate::test_utils::with_connection;
     use xmtp_common::rand_vec;
 
+    #[xmtp_common::test]
+    fn min_key_package_delete_at_ns_none_when_empty() {
+        with_connection(|conn| {
+            // Aggregate MIN over an empty/unmarked table is NULL -> None.
+            assert_eq!(conn.min_key_package_delete_at_ns().unwrap(), None);
+        })
+    }
+
     #[xmtp_common::test]
     fn test_store_key_package_history_entry() {
         with_connection(|conn| {

crates/xmtp_db/src/encrypted_store/tasks.rs

@@ -131,6 +131,10 @@ pub fn data_hash_for(task: &TaskProto) -> TaskDataHash {
     TaskDataHash(xmtp_common::sha256_array(&bytes))
 }
 
+/// Never reaped by expiry: the row's lifetime is bounded by other means (a
+/// recurring row lives forever; an applied pull-in self-deletes).
+pub const NEVER_EXPIRES: i64 = i64::MAX;
+
 pub trait QueryTasks {
     fn create_task(&self, task: NewTask) -> Result<Task, StorageError>;
 

crates/xmtp_db/src/mock.rs

@@ -501,10 +501,13 @@ mock! {
 
     impl QueryIdentity for DbQuery {
         fn queue_key_package_rotation(&self) -> Result<(), StorageError>;
+        fn queue_key_rotation_with_nudge(&self, rotation_task_hash: &crate::tasks::TaskDataHash, rotation_seed: crate::tasks::NewTask) -> Result<(), StorageError>;
 
         fn reset_key_package_rotation_queue(&self, rotation_interval: i64) -> Result<(), StorageError>;
 
         fn is_identity_needs_rotation(&self) -> Result<bool, StorageError>;
+
+        fn next_key_package_rotation_ns(&self) -> Result<Option<i64>, StorageError>;
     }
 
     impl QueryIdentityCache for DbQuery {
@@ -545,6 +548,8 @@ mock! {
             &self,
         ) -> Result<Vec<crate::key_package_history::StoredKeyPackageHistoryEntry>, StorageError>;
 
+        fn min_key_package_delete_at_ns(&self) -> Result<Option<i64>, StorageError>;
+
         fn delete_key_package_history_up_to_id(&self, id: i32) -> Result<(), StorageError>;
 
         fn delete_key_package_entry_with_id(&self, id: i32) -> Result<(), StorageError>;