use std::sync::Arc;
use std::time::SystemTime;
use kmp_domain::{
BundleMetadata, GraphNeighborhoodReader, KmpBundle, NodeDetailReader, SnapshotSaveOptions,
SnapshotStore,
};
use crate::ApplicationError;
use crate::queries::{
ContextRenderOptions, NodeCentricProjectionReader, QueryApplicationService,
QueryTimingBreakdown,
context_render_options::EndpointHint,
render_graph_bundle::{RenderedContext, render_graph_bundle_with_options},
};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RehydrateSessionQuery {
pub root_node_id: String,
pub roles: Vec<String>,
pub persist_snapshot: bool,
pub timeline_window: u32,
pub snapshot_ttl_seconds: u64,
}
#[derive(Debug, Clone, PartialEq)]
pub struct RehydrateSessionResult {
pub root_node_id: String,
pub bundles: Vec<KmpBundle>,
pub rendered_contexts: Vec<RenderedContext>,
pub timeline_events: u32,
pub version: BundleMetadata,
pub snapshot_persisted: bool,
pub snapshot_id: Option<String>,
pub generated_at: SystemTime,
pub timing: Option<QueryTimingBreakdown>,
}
#[derive(Debug)]
pub struct RehydrateSessionUseCase<G, D, S> {
graph_reader: G,
detail_reader: D,
snapshot_store: S,
generator_version: &'static str,
}
impl<G, D, S> RehydrateSessionUseCase<G, D, S>
where
G: GraphNeighborhoodReader + Send + Sync,
D: NodeDetailReader + Send + Sync,
S: SnapshotStore + Send + Sync,
{
pub fn new(
graph_reader: G,
detail_reader: D,
snapshot_store: S,
generator_version: &'static str,
) -> Self {
Self {
graph_reader,
detail_reader,
snapshot_store,
generator_version,
}
}
pub async fn execute(
&self,
root_node_id: &str,
role: &str,
persist_snapshot: bool,
snapshot_options: SnapshotSaveOptions,
) -> Result<(KmpBundle, QueryTimingBreakdown), ApplicationError> {
self.execute_with_depth(
root_node_id,
role,
crate::queries::DEFAULT_NATIVE_GRAPH_TRAVERSAL_DEPTH,
persist_snapshot,
snapshot_options,
)
.await
}
pub async fn execute_with_depth(
&self,
root_node_id: &str,
role: &str,
depth: u32,
persist_snapshot: bool,
snapshot_options: SnapshotSaveOptions,
) -> Result<(KmpBundle, QueryTimingBreakdown), ApplicationError> {
let bundle_reader =
NodeCentricProjectionReader::new(&self.graph_reader, &self.detail_reader);
let (bundle, timing) = match bundle_reader
.load_bundle_with_depth(root_node_id, role, self.generator_version, depth)
.await?
{
(Some(bundle), timing) => (bundle, timing),
(None, _) => {
return Err(ApplicationError::NotFound(format!(
"node '{}' not found",
root_node_id
)));
}
};
if persist_snapshot {
self.snapshot_store
.save_bundle_with_options(&bundle, snapshot_options)
.await?;
}
Ok((bundle, timing))
}
}
impl<G, D, S> QueryApplicationService<G, D, S>
where
G: GraphNeighborhoodReader + Send + Sync,
D: NodeDetailReader + Send + Sync,
S: SnapshotStore + Send + Sync,
{
pub async fn rehydrate_session(
&self,
query: RehydrateSessionQuery,
) -> Result<RehydrateSessionResult, ApplicationError> {
if query.roles.is_empty() {
return Err(ApplicationError::Validation(
"roles cannot be empty".to_string(),
));
}
let bundle_reader = NodeCentricProjectionReader::new(
self.graph_reader.as_ref(),
self.detail_reader.as_ref(),
);
let snapshot_options = SnapshotSaveOptions::new(Some(query.snapshot_ttl_seconds));
let (bundles, timing) = bundle_reader
.load_bundles_for_roles(
&query.root_node_id,
&query.roles,
self.generator_version,
crate::queries::DEFAULT_NATIVE_GRAPH_TRAVERSAL_DEPTH,
)
.await?;
let bundles = match bundles {
Some(bundles) => bundles,
None => {
return Err(ApplicationError::NotFound(format!(
"node '{}' not found",
query.root_node_id
)));
}
};
let session_options = ContextRenderOptions {
endpoint_hint: EndpointHint::SessionSnapshot,
..Default::default()
};
let rendered_contexts: Vec<RenderedContext> = bundles
.iter()
.map(|b| render_graph_bundle_with_options(b, &session_options))
.collect();
if query.persist_snapshot {
for bundle in &bundles {
self.snapshot_store
.save_bundle_with_options(bundle, snapshot_options)
.await?;
}
}
let snapshot_id = if query.persist_snapshot {
Some(format!(
"snapshot:{}:{}",
query.root_node_id,
query.roles.join(",")
))
} else {
None
};
Ok(RehydrateSessionResult {
root_node_id: query.root_node_id,
bundles,
rendered_contexts,
timeline_events: query.timeline_window,
version: BundleMetadata::initial(self.generator_version),
snapshot_persisted: query.persist_snapshot,
snapshot_id,
generated_at: SystemTime::now(),
timing: Some(timing),
})
}
pub async fn warmup_bundle(&self) -> Result<KmpBundle, ApplicationError> {
let (bundle, _timing) = RehydrateSessionUseCase::new(
Arc::clone(&self.graph_reader),
Arc::clone(&self.detail_reader),
Arc::clone(&self.snapshot_store),
self.generator_version,
)
.execute(
"bootstrap-node",
"system",
false,
SnapshotSaveOptions::default(),
)
.await?;
Ok(bundle)
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use std::sync::Arc;
use tokio::sync::Mutex;
use kmp_domain::{
ContextPathNeighborhood, NodeDetailProjection, NodeNeighborhood, NodeProjection, PortError,
SnapshotSaveOptions,
};
use super::{QueryApplicationService, RehydrateSessionQuery};
struct SeededGraphReader;
impl kmp_domain::GraphNeighborhoodReader for SeededGraphReader {
async fn load_neighborhood(
&self,
root_node_id: &str,
_depth: u32,
) -> Result<Option<NodeNeighborhood>, PortError> {
Ok(Some(NodeNeighborhood {
root: NodeProjection {
node_id: root_node_id.to_string(),
node_kind: "story".to_string(),
title: "Root".to_string(),
summary: "Root summary".to_string(),
status: "ACTIVE".to_string(),
labels: vec!["Story".to_string()],
properties: BTreeMap::new(),
provenance: None,
},
neighbors: Vec::new(),
relations: Vec::new(),
}))
}
async fn load_context_path(
&self,
_root_node_id: &str,
_target_node_id: &str,
_subtree_depth: u32,
) -> Result<Option<ContextPathNeighborhood>, PortError> {
Ok(None)
}
}
struct SeededDetailReader;
impl kmp_domain::NodeDetailReader for SeededDetailReader {
async fn load_node_detail(
&self,
node_id: &str,
) -> Result<Option<NodeDetailProjection>, PortError> {
Ok(Some(NodeDetailProjection {
node_id: node_id.to_string(),
detail: "Expanded detail".to_string(),
content_hash: "hash-1".to_string(),
revision: 2,
}))
}
async fn load_node_details_batch(
&self,
node_ids: Vec<String>,
) -> Result<Vec<Option<NodeDetailProjection>>, PortError> {
let mut results = Vec::with_capacity(node_ids.len());
for node_id in &node_ids {
results.push(self.load_node_detail(node_id).await?);
}
Ok(results)
}
}
#[derive(Debug, Default)]
struct RecordingSnapshotStore {
options: Mutex<Vec<SnapshotSaveOptions>>,
}
impl kmp_domain::SnapshotStore for RecordingSnapshotStore {
async fn save_bundle_with_options(
&self,
_bundle: &kmp_domain::KmpBundle,
options: SnapshotSaveOptions,
) -> Result<(), PortError> {
self.options.lock().await.push(options);
Ok(())
}
}
#[tokio::test]
async fn rehydrate_session_propagates_snapshot_ttl_to_store() {
let snapshot_store = Arc::new(RecordingSnapshotStore::default());
let service = QueryApplicationService::new(
Arc::new(SeededGraphReader),
Arc::new(SeededDetailReader),
Arc::clone(&snapshot_store),
"0.1.0",
);
let result = service
.rehydrate_session(RehydrateSessionQuery {
root_node_id: "story-123".to_string(),
roles: vec!["developer".to_string()],
persist_snapshot: true,
timeline_window: 50,
snapshot_ttl_seconds: 1800,
})
.await
.expect("rehydration should succeed");
assert!(result.snapshot_persisted);
assert_eq!(
snapshot_store.options.lock().await.as_slice(),
&[SnapshotSaveOptions::new(Some(1800))]
);
}
struct EmptyGraphReader;
impl kmp_domain::GraphNeighborhoodReader for EmptyGraphReader {
async fn load_neighborhood(
&self,
_root_node_id: &str,
_depth: u32,
) -> Result<Option<NodeNeighborhood>, PortError> {
Ok(None)
}
async fn load_context_path(
&self,
_root_node_id: &str,
_target_node_id: &str,
_subtree_depth: u32,
) -> Result<Option<ContextPathNeighborhood>, PortError> {
Ok(None)
}
}
#[tokio::test]
async fn rehydrate_session_returns_not_found_when_node_does_not_exist() {
let service = QueryApplicationService::new(
Arc::new(EmptyGraphReader),
Arc::new(SeededDetailReader),
Arc::new(RecordingSnapshotStore::default()),
"0.1.0",
);
let result = service
.rehydrate_session(RehydrateSessionQuery {
root_node_id: "nonexistent-node".to_string(),
roles: vec!["developer".to_string()],
persist_snapshot: false,
timeline_window: 0,
snapshot_ttl_seconds: 0,
})
.await;
match result {
Err(crate::ApplicationError::NotFound(_)) => {}
other => panic!("expected NotFound, got: {other:?}"),
}
}
#[tokio::test]
async fn rehydrate_session_multi_role_returns_one_bundle_per_role() {
let service = QueryApplicationService::new(
Arc::new(SeededGraphReader),
Arc::new(SeededDetailReader),
Arc::new(RecordingSnapshotStore::default()),
"0.1.0",
);
let result = service
.rehydrate_session(RehydrateSessionQuery {
root_node_id: "story-123".to_string(),
roles: vec![
"developer".to_string(),
"reviewer".to_string(),
"ops".to_string(),
],
persist_snapshot: false,
timeline_window: 0,
snapshot_ttl_seconds: 0,
})
.await
.expect("multi-role rehydration should succeed");
assert_eq!(result.bundles.len(), 3);
assert_eq!(result.bundles[0].role().as_str(), "developer");
assert_eq!(result.bundles[1].role().as_str(), "reviewer");
assert_eq!(result.bundles[2].role().as_str(), "ops");
for bundle in &result.bundles {
assert_eq!(bundle.root_node().node_id(), "story-123");
assert_eq!(bundle.node_details().len(), 1);
}
}
#[tokio::test]
async fn rehydrate_session_multi_role_not_found_returns_error_not_partial() {
let service = QueryApplicationService::new(
Arc::new(EmptyGraphReader),
Arc::new(SeededDetailReader),
Arc::new(RecordingSnapshotStore::default()),
"0.1.0",
);
let result = service
.rehydrate_session(RehydrateSessionQuery {
root_node_id: "nonexistent".to_string(),
roles: vec!["developer".to_string(), "reviewer".to_string()],
persist_snapshot: false,
timeline_window: 0,
snapshot_ttl_seconds: 0,
})
.await;
match result {
Err(crate::ApplicationError::NotFound(_)) => {}
other => panic!("expected NotFound for all roles, got: {other:?}"),
}
}
#[tokio::test]
async fn rehydrate_session_multi_role_persists_snapshot_per_bundle() {
let snapshot_store = Arc::new(RecordingSnapshotStore::default());
let service = QueryApplicationService::new(
Arc::new(SeededGraphReader),
Arc::new(SeededDetailReader),
Arc::clone(&snapshot_store),
"0.1.0",
);
let result = service
.rehydrate_session(RehydrateSessionQuery {
root_node_id: "story-123".to_string(),
roles: vec!["developer".to_string(), "reviewer".to_string()],
persist_snapshot: true,
timeline_window: 0,
snapshot_ttl_seconds: 900,
})
.await
.expect("multi-role rehydration should succeed");
assert!(result.snapshot_persisted);
assert_eq!(
result.snapshot_id,
Some("snapshot:story-123:developer,reviewer".to_string())
);
assert_eq!(snapshot_store.options.lock().await.len(), 2);
}
#[tokio::test]
async fn rehydrate_session_empty_roles_returns_validation_error() {
let service = QueryApplicationService::new(
Arc::new(SeededGraphReader),
Arc::new(SeededDetailReader),
Arc::new(RecordingSnapshotStore::default()),
"0.1.0",
);
let result = service
.rehydrate_session(RehydrateSessionQuery {
root_node_id: "story-123".to_string(),
roles: vec![],
persist_snapshot: false,
timeline_window: 0,
snapshot_ttl_seconds: 0,
})
.await;
match result {
Err(crate::ApplicationError::Validation(msg)) => {
assert!(msg.contains("roles"), "error should mention roles: {msg}");
}
other => panic!("expected Validation error, got: {other:?}"),
}
}
}