diff --git a/Cargo.lock b/Cargo.lock index dc49190b..c1f1b409 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -392,7 +392,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -1133,7 +1133,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -1328,7 +1328,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -1652,7 +1652,7 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -1770,9 +1770,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.53.1" +version = "1.53.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +checksum = "ce3335fa71841cda333a58d7615b03901380ecf09d59b3296d21f8bbac0dde4e" dependencies = [ "bytes", "libc", @@ -2182,7 +2182,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] diff --git a/crates/tinymemory-tools/src/lifecycle/turn.rs b/crates/tinymemory-tools/src/lifecycle/turn.rs index d39209b5..2dad6f83 100644 --- a/crates/tinymemory-tools/src/lifecycle/turn.rs +++ b/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; + async fn resolve(&self) -> tinymemory_api::Result; } /// Receives background jobs returned by a successful post-turn log. @@ -307,7 +308,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; } @@ -427,7 +428,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 { @@ -483,7 +484,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; } diff --git a/crates/tinymemory-tools/src/lifecycle/turn_runtime_tests.rs b/crates/tinymemory-tools/src/lifecycle/turn_runtime_tests.rs index dee04197..fb3ce185 100644 --- a/crates/tinymemory-tools/src/lifecycle/turn_runtime_tests.rs +++ b/crates/tinymemory-tools/src/lifecycle/turn_runtime_tests.rs @@ -138,8 +138,27 @@ struct StaticResolver { advance: Duration, } +struct DeferredResolver { + memory: AgentMemory, + release: Mutex>>, + calls: AtomicUsize, +} + +#[async_trait] +impl EngineResolver for DeferredResolver { + async fn resolve(&self) -> Result { + 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 { + async fn resolve(&self) -> Result { 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,