diff --git a/crates/tracedecay-session-temporal-store/src/lib.rs b/crates/tracedecay-session-temporal-store/src/lib.rs index 8812d0673a..e69c647ebf 100644 --- a/crates/tracedecay-session-temporal-store/src/lib.rs +++ b/crates/tracedecay-session-temporal-store/src/lib.rs @@ -94,7 +94,7 @@ use self::hydration::GlobalDbTemporalHydrationPort; use self::participant_freeze::{ freeze_participants, freeze_prepared_candidate_participants, root_readiness, }; -use self::retrieval::{GlobalDbPreparedCandidatePort, GlobalDbTemporalReadPort}; +use self::retrieval::GlobalDbTemporalReadPort; use self::sql::TemporalSqlRead; use tracedecay_lcm::payload::read_verified_payload_content_with_checkpoint; @@ -893,16 +893,8 @@ impl<'db, D: SessionTemporalRegisteredDb + Sync> request.direct_anchor(), request.snapshot_request().semantic_filter().goals, ); - let preparation = GlobalDbPreparedCandidatePort::new( - &candidate_read, - request.snapshot_request(), - &plan, - ); - let prepared = - tracedecay_temporal_query::ports::prepare_temporal_candidate_cohort( - request.snapshot_request(), - &preparation, - ) + let prepared = candidate_read + .prepare_root_candidate_cohort(request.snapshot_request(), &plan) .await .map_err(map_control_error)?; let readiness = root_readiness(&temporal_read, request).await?; diff --git a/crates/tracedecay-session-temporal-store/src/retrieval.rs b/crates/tracedecay-session-temporal-store/src/retrieval.rs index eec0edffe1..ea59cff40e 100644 --- a/crates/tracedecay-session-temporal-store/src/retrieval.rs +++ b/crates/tracedecay-session-temporal-store/src/retrieval.rs @@ -22,10 +22,12 @@ use tracedecay_domain::{ use tracedecay_runtime_core::db::engine; use tracedecay_temporal_query::candidates::{CandidateChannel, CandidatePlan}; use tracedecay_temporal_query::ports::{ - CandidatePageSink, MeasuredTemporalValue, PageRequest, PageStatus, PortFuture, - TemporalCandidateFilterV1, TemporalCandidatePreparationPort, TemporalExecutionSnapshot, - TemporalMessageTypeFilterV1, TemporalPortError, TemporalReadPort, TemporalRecordPageSink, + CANDIDATE_READ_BUDGET, CandidateFieldCaps, CandidatePageSink, CandidateReadState, + MeasuredTemporalValue, PageLimits, PageRequest, PageStatus, PortFuture, + TemporalCandidateFilterV1, TemporalExecutionSnapshot, TemporalMessageTypeFilterV1, + TemporalPortError, TemporalPreparedCandidateCohort, TemporalReadPort, TemporalRecordPageSink, TemporalRetrievalScope, TemporalSessionScopeFilterV1, TemporalSnapshotRequest, + await_controlled, begin_prepared_candidate_pull, commit_prepared_candidate_pull, }; use tracedecay_temporal_query::ranking::RankingCandidate; @@ -182,48 +184,6 @@ pub struct GlobalDbTemporalReadPort<'a> { git_scope_session_ids: Option<&'a BTreeSet<(String, String)>>, } -pub(super) struct GlobalDbPreparedCandidatePort<'port, 'db, 'request> { - read_port: &'port GlobalDbTemporalReadPort<'db>, - request: &'request TemporalSnapshotRequest, - plan: &'request CandidatePlan, -} - -impl<'port, 'db, 'request> GlobalDbPreparedCandidatePort<'port, 'db, 'request> { - #[hotpath::skip] - pub(super) const fn new( - read_port: &'port GlobalDbTemporalReadPort<'db>, - request: &'request TemporalSnapshotRequest, - plan: &'request CandidatePlan, - ) -> Self { - Self { - read_port, - request, - plan, - } - } -} - -impl TemporalCandidatePreparationPort for GlobalDbPreparedCandidatePort<'_, '_, '_> { - fn produce_prepared_candidate_page<'a>( - &'a self, - request: PageRequest, - sink: &'a mut CandidatePageSink<'_>, - ) -> PortFuture<'a, PageStatus> { - Box::pin(async move { - self.read_port - .produce_candidates_from_request( - &TemporalRetrievalScope::AllSessionsInAuthorizedRoot, - self.request, - 1, - self.plan, - &request, - sink, - ) - .await - }) - } -} - struct SessionReadRelationAuthority<'a> { scope: &'a SessionRelationScope, store: SessionRelationGraphStore, @@ -276,6 +236,65 @@ impl<'a> GlobalDbTemporalReadPort<'a> { } } + #[hotpath::measure( + future = true, + label = "session_temporal.query.prepare_root_candidates" + )] + pub(super) async fn prepare_root_candidate_cohort( + &self, + request: &TemporalSnapshotRequest, + plan: &CandidatePlan, + ) -> Result { + request.execution_control().checkpoint()?; + let limits = request.limits(); + let candidate_page_items = limits.candidate_limit.min(64); + let candidate_limits = PageLimits::new( + limits.candidate_limit, + limits.candidate_total_bytes, + limits.candidate_item_bytes, + candidate_page_items, + )?; + let mut state = CandidateReadState::new(candidate_limits); + let mut candidates = Vec::with_capacity(limits.candidate_limit.min(256)); + let scope = TemporalRetrievalScope::AllSessionsInAuthorizedRoot; + loop { + let limits = begin_prepared_candidate_pull(request, &mut state)?; + let control = request.execution_control(); + let field_caps = CandidateFieldCaps::new( + limits.candidate_stable_id_bytes, + limits.candidate_anchor_id_bytes, + limits.candidate_metadata_field_bytes, + ); + let page_request = state.request(limits.candidate_key_bytes, Some(field_caps)); + let mut sink = state.begin_page( + control, + limits.candidate_key_bytes, + Some(field_caps), + CANDIDATE_READ_BUDGET, + ); + let status = await_controlled( + control, + Box::pin(self.produce_candidates_from_request( + &scope, + request, + 1, + plan, + &page_request, + &mut sink, + )), + ) + .await?; + let page = sink.finish(status)?; + let page = commit_prepared_candidate_pull(&mut state, page)?; + let status = page.status(); + candidates.extend(page.into_items()); + if status == PageStatus::Complete { + break; + } + } + TemporalPreparedCandidateCohort::new(candidates) + } + #[hotpath::skip] async fn candidate_matches_filter( &self, diff --git a/crates/tracedecay-session-temporal-store/src/retrieval/tests.rs b/crates/tracedecay-session-temporal-store/src/retrieval/tests.rs index 0eafa07b06..c7ed01dd3e 100644 --- a/crates/tracedecay-session-temporal-store/src/retrieval/tests.rs +++ b/crates/tracedecay-session-temporal-store/src/retrieval/tests.rs @@ -15,12 +15,15 @@ use tracedecay_runtime_core::db::{ engine::{Connection, Executor, TestConnection, Value as SqlValue}, }; use tracedecay_temporal_query::candidates::CandidateChannel; +use tracedecay_temporal_query::plan_temporal_candidates; use tracedecay_temporal_query::ports::{ - BindingDigest, ExecutionControl, KernelVersions, PageRequest, TemporalAuthorizedRoot, + BindingDigest, CANDIDATE_READ_BUDGET, CandidateFieldCaps, CandidateReadState, ExecutionControl, + ExecutionLimits, KernelVersions, PageLimits, PageRequest, PageStatus, TemporalAuthorizedRoot, TemporalExecutionSnapshot, TemporalParticipantAuthorization, TemporalParticipantGeneration, TemporalParticipantManifest, TemporalPortError, TemporalPreparedCandidateCohort, TemporalRecord, TemporalRetrievalScope, TemporalSnapshotRequest, TemporalSourceAccess, - TemporalWatermarks, + TemporalWatermarks, await_controlled, begin_prepared_candidate_pull, + commit_prepared_candidate_pull, }; use tracedecay_temporal_query::ranking::RankingCandidate; use tracedecay_temporal_query::resolution::{SummarySourceState, ValidatedAuthorization}; @@ -234,6 +237,123 @@ fn prepared_root_candidates_bind_each_frozen_participant_generation() { ); } +fn root_preparation_request(mode: TemporalModeV1) -> TemporalSnapshotRequest { + root_snapshot_with_mode(1, None, mode).request().clone() +} + +#[tokio::test] +async fn root_candidate_preparation_matches_direct_request_producer() { + let dir = tempdir().expect("temporary directory"); + let runtime = HostAdmissionTestRuntimeV1::profile(dir.path()) + .await + .expect("registered profile runtime"); + runtime.seed_candidate_query_fixture_for_test().await; + let read = runtime.retrieval_read_for_test().await; + let adapter = read.adapter(); + let request = root_preparation_request(TemporalModeV1::Current); + let plan = plan_temporal_candidates("needle candidate", None, false); + let direct = adapter + .prepare_root_candidate_cohort(&request, &plan) + .await + .expect("direct root candidate cohort"); + let candidate_limits = PageLimits::new( + request.limits().candidate_limit, + request.limits().candidate_total_bytes, + request.limits().candidate_item_bytes, + request.limits().candidate_limit.min(64), + ) + .expect("valid limits"); + let mut state = CandidateReadState::new(candidate_limits); + let mut via_pages = Vec::new(); + let scope = TemporalRetrievalScope::AllSessionsInAuthorizedRoot; + loop { + let limits = begin_prepared_candidate_pull(&request, &mut state).expect("pull limits"); + let control = request.execution_control(); + let field_caps = CandidateFieldCaps::new( + limits.candidate_stable_id_bytes, + limits.candidate_anchor_id_bytes, + limits.candidate_metadata_field_bytes, + ); + let page_request = state.request(limits.candidate_key_bytes, Some(field_caps)); + let mut sink = state.begin_page( + control, + limits.candidate_key_bytes, + Some(field_caps), + CANDIDATE_READ_BUDGET, + ); + let status = await_controlled( + control, + Box::pin(adapter.produce_candidates_from_request( + &scope, + &request, + 1, + &plan, + &page_request, + &mut sink, + )), + ) + .await + .expect("page-driven root candidate cohort"); + let page = sink.finish(status).expect("bounded page"); + let page = commit_prepared_candidate_pull(&mut state, page).expect("committed page"); + let status = page.status(); + via_pages.extend(page.into_items()); + if status == PageStatus::Complete { + break; + } + } + assert_eq!( + direct.candidates(), + via_pages.as_slice(), + "direct preparation must match the request producer path" + ); + assert!( + !direct.candidates().is_empty(), + "root candidate preparation must return live fixture candidates" + ); +} + +#[tokio::test] +async fn root_candidate_preparation_preserves_item_byte_budget_failure() { + let dir = tempdir().expect("temporary directory"); + let runtime = HostAdmissionTestRuntimeV1::profile(dir.path()) + .await + .expect("registered profile runtime"); + runtime.seed_candidate_query_fixture_for_test().await; + let read = runtime.retrieval_read_for_test().await; + let adapter = read.adapter(); + let request = root_preparation_request(TemporalModeV1::Current).with_limits(ExecutionLimits { + candidate_limit: 8, + candidate_total_bytes: 64 * 1024, + candidate_item_bytes: 8, + ..ExecutionLimits::default() + }); + let plan = plan_temporal_candidates("needle candidate", None, false); + assert!(matches!( + adapter.prepare_root_candidate_cohort(&request, &plan).await, + Err(TemporalPortError::BudgetExceeded { .. }) + )); +} + +#[tokio::test] +async fn root_candidate_preparation_preserves_live_cancellation() { + let dir = tempdir().expect("temporary directory"); + let runtime = HostAdmissionTestRuntimeV1::profile(dir.path()) + .await + .expect("registered profile runtime"); + runtime.seed_candidate_query_fixture_for_test().await; + let read = runtime.retrieval_read_for_test().await; + let adapter = read.adapter(); + let control = ExecutionControl::default(); + control.cancel(); + let request = root_preparation_request(TemporalModeV1::Current).with_execution_control(control); + let plan = plan_temporal_candidates("needle candidate", None, false); + assert_eq!( + adapter.prepare_root_candidate_cohort(&request, &plan).await, + Err(TemporalPortError::Cancelled) + ); +} + fn record_request() -> PageRequest { PageRequest::for_test(32, 64 * 1024, 8 * 1024, 32, 512) } diff --git a/crates/tracedecay-temporal-query/src/ports/contracts.rs b/crates/tracedecay-temporal-query/src/ports/contracts.rs index 93aa887f65..8050ca4c3a 100644 --- a/crates/tracedecay-temporal-query/src/ports/contracts.rs +++ b/crates/tracedecay-temporal-query/src/ports/contracts.rs @@ -181,6 +181,31 @@ pub async fn prepare_temporal_candidate_cohort( TemporalPreparedCandidateCohort::new(candidates) } +pub fn begin_prepared_candidate_pull( + request: &TemporalSnapshotRequest, + state: &mut CandidateReadState, +) -> Result { + begin_pull_request( + request, + state, + |limits| { + ( + limits.candidate_limit, + limits.candidate_total_bytes, + limits.candidate_item_bytes, + ) + }, + CANDIDATE_READ_BUDGET, + ) +} + +pub fn commit_prepared_candidate_pull( + state: &mut CandidateReadState, + page: BoundedPage, +) -> Result, TemporalPortError> { + commit_pulled_page(state, page, CANDIDATE_READ_BUDGET) +} + pub async fn pull_candidate_page( port: &impl TemporalReadPort, snapshot: &TemporalExecutionSnapshot, diff --git a/crates/tracedecay-temporal-query/src/ports/execution.rs b/crates/tracedecay-temporal-query/src/ports/execution.rs index 852893a7a4..b2ddec6a65 100644 --- a/crates/tracedecay-temporal-query/src/ports/execution.rs +++ b/crates/tracedecay-temporal-query/src/ports/execution.rs @@ -238,7 +238,7 @@ impl ExecutionControl { } } -pub(crate) async fn await_controlled( +pub async fn await_controlled( control: &ExecutionControl, future: impl Future>, ) -> Result diff --git a/crates/tracedecay-temporal-query/src/ports/paging.rs b/crates/tracedecay-temporal-query/src/ports/paging.rs index 25152a3309..e91660c7b5 100644 --- a/crates/tracedecay-temporal-query/src/ports/paging.rs +++ b/crates/tracedecay-temporal-query/src/ports/paging.rs @@ -64,7 +64,7 @@ pub struct CandidateFieldCaps { impl CandidateFieldCaps { #[hotpath::skip] - pub(super) const fn new( + pub const fn new( stable_id_bytes: usize, anchor_id_bytes: usize, metadata_field_bytes: usize, @@ -256,7 +256,7 @@ impl ReadState { self.consumed_bytes } - pub(super) fn require_within_limits( + pub fn require_within_limits( &self, max_items: usize, max_total_bytes: usize, @@ -281,7 +281,7 @@ impl ReadState { Ok(()) } - pub(super) fn request( + pub fn request( &self, max_key_bytes: usize, candidate_field_caps: Option, @@ -303,12 +303,12 @@ impl ReadState { } } - pub(super) fn is_exhausted(&self) -> bool { + pub fn is_exhausted(&self) -> bool { self.consumed_items == self.limits.max_items || self.consumed_bytes == self.limits.max_total_bytes } - pub(super) fn begin_page<'a>( + pub fn begin_page<'a>( &'a mut self, control: &'a ExecutionControl, max_key_bytes: usize, @@ -332,15 +332,12 @@ impl ReadState { } } - pub(super) fn advanced_page(&mut self, continuation: Option) { + pub fn advanced_page(&mut self, continuation: Option) { self.page_index += 1; self.keyset = continuation; } - pub(super) fn incomplete_coverage_error( - &self, - resources: ReadBudgetResources, - ) -> TemporalPortError { + pub fn incomplete_coverage_error(&self, resources: ReadBudgetResources) -> TemporalPortError { if self.consumed_items == self.limits.max_items { TemporalPortError::BudgetExceeded { resource: resources.item_count, @@ -354,13 +351,13 @@ impl ReadState { } #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(super) struct ReadBudgetResources { +pub struct ReadBudgetResources { item_count: &'static str, item_bytes: &'static str, total_bytes: &'static str, } -pub(super) const CANDIDATE_READ_BUDGET: ReadBudgetResources = ReadBudgetResources { +pub const CANDIDATE_READ_BUDGET: ReadBudgetResources = ReadBudgetResources { item_count: "candidate item count", item_bytes: "candidate item bytes", total_bytes: "candidate total bytes", @@ -448,7 +445,7 @@ impl BoundedPageSink<'_, T> { Ok(()) } - pub(super) fn finish(self, status: PageStatus) -> Result, TemporalPortError> { + pub fn finish(self, status: PageStatus) -> Result, TemporalPortError> { if status == PageStatus::More && self.items.is_empty() { return Err(TemporalPortError::Read { operation: "produce bounded page",