chore: roll up security migration changes

crates/tinymemory-tools/src/deletion/mod.rs

@@ -88,7 +88,7 @@ impl Backoff {
 pub trait DeletionExecutor: Send + Sync {
     /// Whether the engine can be reached at all. When not, a drain stops
     /// before running anything.
-    fn is_reachable(&self) -> bool;
+    async fn is_reachable(&self) -> bool;
 
     /// Runs `deletion`, returning how many items went.
     ///
@@ -252,7 +252,7 @@ impl<S: ListStore<PendingDeletion>> DeletionQueue<S> {
         if queued.is_empty() {
             return 0;
         }
-        if !executor.is_reachable() {
+        if !executor.is_reachable().await {
             log::debug!(
                 "[tinymemory:deletion] memory off; {} pending deletions wait",
                 queued.len()

crates/tinymemory-tools/src/deletion/mod_tests.rs

@@ -1,8 +1,12 @@
 use super::*;
 
+use std::future::Future;
 use std::sync::Mutex as StdMutex;
+use std::task::Poll;
 
 use crate::list_store::MemoryListStore;
+use futures::channel::oneshot;
+use futures::future::poll_fn;
 
 fn thread(id: &str) -> PendingDeletion {
     PendingDeletion::Thread {
@@ -22,7 +26,7 @@ struct Engine {
 
 #[async_trait]
 impl DeletionExecutor for Engine {
-    fn is_reachable(&self) -> bool {
+    async fn is_reachable(&self) -> bool {
         self.reachable
     }
     async fn run(&self, deletion: &PendingDeletion) -> Result<usize, String> {
@@ -35,6 +39,27 @@ impl DeletionExecutor for Engine {
     }
 }
 
+struct DeferredReachability {
+    release: StdMutex<Option<oneshot::Receiver<()>>>,
+    ran: StdMutex<Vec<PendingDeletion>>,
+}
+
+#[async_trait]
+impl DeletionExecutor for DeferredReachability {
+    async fn is_reachable(&self) -> bool {
+        let release = self.release.lock().unwrap().take();
+        if let Some(release) = release {
+            let _ = release.await;
+        }
+        true
+    }
+
+    async fn run(&self, deletion: &PendingDeletion) -> Result<usize, String> {
+        self.ran.lock().unwrap().push(deletion.clone());
+        Ok(1)
+    }
+}
+
 fn engine(reachable: bool, failing: bool) -> Engine {
     Engine {
         reachable,
@@ -64,6 +89,29 @@ async fn nothing_runs_while_the_engine_is_unreachable() {
     assert_eq!(queue.pending(), vec![thread("t")]);
 }
 
+#[tokio::test]
+async fn drain_awaits_async_reachability_before_running_deletions() {
+    let queue = queue();
+    queue.enqueue(thread("t"));
+    let (release, wait) = oneshot::channel();
+    let engine = DeferredReachability {
+        release: StdMutex::new(Some(wait)),
+        ran: StdMutex::new(Vec::new()),
+    };
+    let mut drain = Box::pin(queue.drain(&engine, Drain::Force, Utc::now()));
+
+    poll_fn(|cx| match drain.as_mut().poll(cx) {
+        Poll::Pending => Poll::Ready(()),
+        Poll::Ready(_) => panic!("drain completed before async reachability resolved"),
+    })
+    .await;
+    assert!(engine.ran.lock().unwrap().is_empty());
+
+    release.send(()).unwrap();
+    assert_eq!(drain.await, 1);
+    assert_eq!(engine.ran.lock().unwrap().as_slice(), &[thread("t")]);
+}
+
 #[tokio::test]
 async fn a_reachable_engine_settles_the_queue() {
     let queue = queue();

crates/tinymemory-tools/src/layout_migration/job.rs

@@ -94,7 +94,7 @@ pub trait LayoutHost: Send + Sync {
     /// # Errors
     ///
     /// When memory is off or the account has no per-user root.
-    fn engines(&self) -> Result<Engines>;
+    async fn engines(&self) -> Result<Engines>;
 
     /// Where items go in the per-user tree.
     ///
@@ -128,7 +128,7 @@ pub trait LayoutHost: Send + Sync {
     /// # Errors
     ///
     /// When memory is off or the account has no per-user root.
-    fn legacy_claim(&self) -> Result<Option<Box<dyn LegacyClaim>>>;
+    async fn legacy_claim(&self) -> Result<Option<Box<dyn LegacyClaim>>>;
 
     /// Whether the last import of old local memory ended with its last batch
     /// not confirmed listed. The migration then forgets by id, never erasing
@@ -167,15 +167,15 @@ where
     if state.phase == Phase::Cleaned && (state.rechecked || !host.listed_unconfirmed()) {
         return Ok(Outcome::Done);
     }
-    let engines = host.engines()?;
+    let engines = host.engines().await?;
     if state.phase == Phase::Idle && !state.switched && !legacy_present(&*engines.legacy).await? {
         return nothing_to_move(host, &mut state).await;
     }
     if let Trigger::Manual { takeover: true } = trigger {
         state.takeover = true;
         store.save(&state)?;
     }
-    let shared = host.legacy_claim()?;
+    let shared = host.legacy_claim().await?;
     if let Some(claim) = &shared {
         if claim.held_by_other()? {
             return claimed_elsewhere(host, &mut state).await;

crates/tinymemory-tools/src/layout_migration/service.rs

@@ -197,7 +197,7 @@ pub struct MigrationStatus {
 /// tree cannot be read.
 pub async fn scan(host: &dyn LayoutHost) -> Result<ScanView> {
     let state = host.store().load()?;
-    let claim = match host.legacy_claim() {
+    let claim = match host.legacy_claim().await {
         // Memory off or signed out: nothing to show.
         Err(Error::Off(_)) => None,
         claim => claim?,
@@ -215,7 +215,7 @@ pub async fn scan(host: &dyn LayoutHost) -> Result<ScanView> {
         {
             false
         }
-        Phase::Idle if !state.switched => match host.engines() {
+        Phase::Idle if !state.switched => match host.engines().await {
             // Memory off or signed out: nothing to show.
             Err(Error::Off(_)) => false,
             Err(error) => return Err(error),

crates/tinymemory-tools/src/layout_migration/service_tests.rs

@@ -208,7 +208,7 @@ impl LayoutHost for CountingHost {
         &self.store
     }
 
-    fn engines(&self) -> crate::layout_migration::Result<Engines> {
+    async fn engines(&self) -> crate::layout_migration::Result<Engines> {
         if self.fail_engines {
             Err(crate::layout_migration::Error::Off(
                 "engine unavailable".to_string(),
@@ -219,7 +219,7 @@ impl LayoutHost for CountingHost {
                 tree: self.inner.tree.clone(),
             })
         } else {
-            self.inner.engines()
+            self.inner.engines().await
         }
     }
 
@@ -239,8 +239,8 @@ impl LayoutHost for CountingHost {
         self.inner.free_now().await
     }
 
-    fn legacy_claim(&self) -> crate::layout_migration::Result<Option<Box<dyn LegacyClaim>>> {
-        self.inner.legacy_claim()
+    async fn legacy_claim(&self) -> crate::layout_migration::Result<Option<Box<dyn LegacyClaim>>> {
+        self.inner.legacy_claim().await
     }
 
     fn listed_unconfirmed(&self) -> bool {

crates/tinymemory-tools/src/layout_migration/test_host.rs

@@ -100,7 +100,7 @@ impl LayoutHost for FakeHost {
         &self.store
     }
 
-    fn engines(&self) -> Result<Engines> {
+    async fn engines(&self) -> Result<Engines> {
         Ok(Engines {
             legacy: self.legacy.clone(),
             tree: self.tree.clone(),
@@ -134,7 +134,7 @@ impl LayoutHost for FakeHost {
         self.free.load(Ordering::SeqCst) && left > 0
     }
 
-    fn legacy_claim(&self) -> Result<Option<Box<dyn LegacyClaim>>> {
+    async fn legacy_claim(&self) -> Result<Option<Box<dyn LegacyClaim>>> {
         Ok(self
             .claim
             .clone()

crates/tinymemory-tools/src/lifecycle/turn.rs

@@ -56,13 +56,14 @@ pub trait Clock: Send + Sync {
 }
 
 /// Resolves the configured engine and identity to their lifecycle instance.
+#[async_trait]
 pub trait EngineResolver: Send + Sync {
     /// Returns the configured memory instance for this host hook.
     ///
     /// # Errors
     ///
     /// Returns the host's engine or identity resolution failure.
-    fn resolve(&self) -> tinymemory_api::Result<AgentMemory>;
+    async fn resolve(&self) -> tinymemory_api::Result<AgentMemory>;
 }
 
 /// Receives background jobs returned by a successful post-turn log.
@@ -310,7 +311,7 @@ impl TurnHooks {
         let work_clock = clock.clone();
         let model = self.date_hint_model.clone();
         let operation = async move {
-            let mut memory = resolver.resolve().ok()?;
+            let mut memory = resolver.resolve().await.ok()?;
             if deadline.is_some_and(|deadline| work_clock.now() >= deadline) {
                 return None;
             }
@@ -430,7 +431,7 @@ impl TurnHooks {
         if !logging {
             return;
         }
-        let Ok(memory) = self.resolver.resolve() else {
+        let Ok(memory) = self.resolver.resolve().await else {
             return;
         };
         let Ok(report) = memory.post_turn(turn).await else {
@@ -486,7 +487,7 @@ impl TurnHooks {
         let resolver = self.resolver.clone();
         let clock = self.clock.clone();
         let operation = async move {
-            let memory = resolver.resolve().ok()?;
+            let memory = resolver.resolve().await.ok()?;
             if deadline.is_some_and(|deadline| clock.now() >= deadline) {
                 return None;
             }

crates/tinymemory-tools/src/lifecycle/turn_runtime_tests.rs

@@ -138,8 +138,27 @@ struct StaticResolver {
     advance: Duration,
 }
 
+struct DeferredResolver {
+    memory: AgentMemory,
+    release: Mutex<Option<oneshot::Receiver<()>>>,
+    calls: AtomicUsize,
+}
+
+#[async_trait]
+impl EngineResolver for DeferredResolver {
+    async fn resolve(&self) -> Result<AgentMemory> {
+        self.calls.fetch_add(1, Ordering::SeqCst);
+        let release = self.release.lock().unwrap().take();
+        if let Some(release) = release {
+            let _ = release.await;
+        }
+        Ok(self.memory.clone())
+    }
+}
+
+#[async_trait]
 impl EngineResolver for StaticResolver {
-    fn resolve(&self) -> Result<AgentMemory> {
+    async fn resolve(&self) -> Result<AgentMemory> {
         self.calls.fetch_add(1, Ordering::SeqCst);
         self.clock.advance_by(self.advance);
         if self.fail {
@@ -296,6 +315,41 @@ fn fixture_with_memory_and_advance(
     }
 }
 
+#[tokio::test]
+async fn pre_turn_awaits_async_engine_resolution() {
+    let engine = Arc::new(ReferenceEngine::new());
+    let (release, wait) = oneshot::channel();
+    let resolver = Arc::new(DeferredResolver {
+        memory: memory(engine),
+        release: Mutex::new(Some(wait)),
+        calls: AtomicUsize::new(0),
+    });
+    let clock = FakeClock::at(Duration::ZERO);
+    let citations = Arc::new(RecordingCitations::default());
+    let hooks = TurnHooks::new(
+        resolver.clone(),
+        Arc::new(clock.clone()),
+        None,
+        Arc::new(RecordingJobs::default()),
+        citations.clone(),
+    );
+    let mut work = Box::pin(hooks.pre_turn(
+        PreTurn::new("thread", 0, "remember this"),
+        options(true, true),
+    ));
+
+    poll_fn(|cx| match work.as_mut().poll(cx) {
+        Poll::Pending => Poll::Ready(()),
+        Poll::Ready(_) => panic!("pre-turn completed before async engine resolution"),
+    })
+    .await;
+    assert_eq!(resolver.calls.load(Ordering::SeqCst), 1);
+    release.send(()).unwrap();
+
+    let _ = work.await;
+    assert_eq!(citations.calls.load(Ordering::SeqCst), 0);
+}
+
 fn options(logging: bool, recall: bool) -> PreTurnOptions {
     PreTurnOptions {
         logging,