diff --git a/crates/tinyagents-session/src/port/memory/mod.rs b/crates/tinyagents-session/src/port/memory/mod.rs index 97063030b..d36262cd9 100644 --- a/crates/tinyagents-session/src/port/memory/mod.rs +++ b/crates/tinyagents-session/src/port/memory/mod.rs @@ -336,6 +336,26 @@ impl TranscriptLocator for InMemoryTranscriptLocator { .map(|(_, history)| history as Arc) } + /// An exact lookup that never registers a stem. The trait default probes + /// through [`Self::open_stem`], which would leave an empty handle behind + /// for every generation asked about, and `begin_generation` would then + /// find its own successor "already exists". Only a written transcript + /// counts, as in the persistent locators. + fn read_session_transcript(&self, session: &SessionRef) -> Option> { + let stem = session_stem(session); + let history = self + .stems + .lock() + .unwrap_or_else(PoisonError::into_inner) + .iter() + .find(|(known, _, _)| known == &stem) + .map(|(_, _, history)| history.clone())?; + match history.read_session() { + Ok(Some(_)) => Some(history as Arc), + _ => None, + } + } + fn open_stem( &self, stem: &str, diff --git a/crates/tinyagents-session/src/port/memory/mod_tests.rs b/crates/tinyagents-session/src/port/memory/mod_tests.rs index 8b60fac65..925b89584 100644 --- a/crates/tinyagents-session/src/port/memory/mod_tests.rs +++ b/crates/tinyagents-session/src/port/memory/mod_tests.rs @@ -271,6 +271,41 @@ fn an_unwritten_successor_is_reserved_until_written_or_dropped() { ); } +#[test] +fn probing_for_the_head_generation_does_not_reserve_the_next_one() { + let locator = InMemoryTranscriptLocator::new("test"); + let session = SessionRef::scoped("thread-probe", "planner"); + let meta = test_meta("thread-probe"); + let baseline = vec![TranscriptMessage::user("hello")]; + locator + .open_session(&session, meta.clone()) + .unwrap() + .append(baseline[0].clone()) + .unwrap(); + + // Resuming asks for the head, which probes the successor generation. + assert_eq!(locator.head_generation(&session), session); + assert!(!locator.session_exists(&session.next_generation())); + assert!(locator.session_chain(&session).len() == 1); + + let (successor, _) = locator + .begin_generation_from_baseline(&session, meta, &baseline) + .expect("a probe leaves no state for the next generation to collide with"); + assert_eq!(successor, session.next_generation()); +} + +#[test] +fn a_begun_generation_stays_reserved_after_probing_for_it() { + let locator = InMemoryTranscriptLocator::new("test"); + let session = SessionRef::scoped("thread-probe-reserved", "planner"); + let meta = test_meta("thread-probe-reserved"); + let (successor, _) = locator.begin_generation(&session, meta.clone()).unwrap(); + + assert!(!locator.session_exists(&successor)); + assert!(locator.read_session_transcript(&successor).is_none()); + assert!(locator.begin_generation(&session, meta).is_err()); +} + #[test] fn completed_retention_uses_completion_time() { let turns = InMemoryTurnStates::default();