kmp-application 0.18.2

use kmp_domain::{
    GraphNeighborhoodReader, KmpBundle, NeighborhoodRequest, NodeDetailReader, SnapshotSaveOptions,
    SnapshotStore,
};

use crate::ApplicationError;
pub use crate::queries::render_graph_bundle::RenderedContext;
use crate::queries::{
    ContextRenderOptions, QueryApplicationService, QueryTimingBreakdown, RehydrateSessionUseCase,
    clamp_native_graph_traversal_depth, render_graph_bundle_with_options,
};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GetContextQuery {
    pub root_node_id: String,
    pub role: String,
    pub depth: u32,
    pub requested_scopes: Vec<String>,
    pub render_options: ContextRenderOptions,
}

#[derive(Debug, Clone, PartialEq)]
pub struct GetContextResult {
    pub bundle: KmpBundle,
    pub rendered: RenderedContext,
    pub requested_scopes: Vec<String>,
    pub served_at: std::time::SystemTime,
    pub timing: Option<QueryTimingBreakdown>,
}

#[derive(Debug)]
pub struct GetContextUseCase<G, D, S> {
    rehydrate_session: RehydrateSessionUseCase<G, D, S>,
}

impl<G, D, S> GetContextUseCase<G, D, S>
where
    G: GraphNeighborhoodReader + Send + Sync,
    D: NodeDetailReader + Send + Sync,
    S: SnapshotStore + Send + Sync,
{
    pub fn new(rehydrate_session: RehydrateSessionUseCase<G, D, S>) -> Self {
        Self { rehydrate_session }
    }

    pub async fn execute(
        &self,
        root_node_id: &str,
        role: &str,
        depth: u32,
        requested_scopes: &[String],
        render_options: &ContextRenderOptions,
    ) -> Result<GetContextResult, ApplicationError> {
        // The scopes were resolved before this query was built. Handing them to
        // the reader is what lets a store narrow on the axis rather than the
        // caller discarding what it should never have loaded.
        let (bundle, timing) = self
            .rehydrate_session
            .execute_for(
                &NeighborhoodRequest::new(root_node_id, clamp_native_graph_traversal_depth(depth))
                    .with_scopes(requested_scopes.to_vec()),
                role,
                false,
                SnapshotSaveOptions::default(),
            )
            .await?;
        let rendered = render_graph_bundle_with_options(&bundle, render_options);

        Ok(GetContextResult {
            bundle,
            rendered,
            requested_scopes: requested_scopes.to_vec(),
            served_at: std::time::SystemTime::now(),
            timing: Some(timing),
        })
    }
}

impl<G, D, S> QueryApplicationService<G, D, S>
where
    G: GraphNeighborhoodReader + Send + Sync,
    D: NodeDetailReader + Send + Sync,
    S: SnapshotStore + Send + Sync,
{
    /// Read the structured context for consumers that do their own projection.
    /// No prompt rendering, tokenization or snapshot write is needed here.
    pub(crate) async fn read_context_bundle(
        &self,
        request: &NeighborhoodRequest,
        role: &str,
    ) -> Result<KmpBundle, ApplicationError> {
        let rehydrate = RehydrateSessionUseCase::new(
            std::sync::Arc::clone(&self.graph_reader),
            std::sync::Arc::clone(&self.detail_reader),
            std::sync::Arc::clone(&self.snapshot_store),
            self.generator_version,
        );
        let (bundle, _) = rehydrate
            .execute_for(request, role, false, SnapshotSaveOptions::default())
            .await?;
        Ok(bundle)
    }

    /// Read graph metadata first; the caller selects canonical bodies within
    /// this same operation snapshot before returning a materialized result.
    pub(crate) async fn read_context_catalogue(
        &self,
        request: &NeighborhoodRequest,
        role: &str,
    ) -> Result<KmpBundle, ApplicationError> {
        Ok(self
            .read_context_catalogue_with_timing(request, role)
            .await?
            .0)
    }

    pub(crate) async fn read_context_catalogue_with_timing(
        &self,
        request: &NeighborhoodRequest,
        role: &str,
    ) -> Result<(KmpBundle, QueryTimingBreakdown), ApplicationError> {
        let reader = crate::queries::NodeCentricProjectionReader::new(
            std::sync::Arc::clone(&self.graph_reader),
            std::sync::Arc::clone(&self.detail_reader),
        );
        // Lane hints cannot remove other coordinates before whole-entry
        // predicates run. Root selection already restricts the about scope.
        // The full reader also retains this catalogue; defer bodies only.
        let catalogue = NeighborhoodRequest::new(request.root_node_id(), request.depth());
        let (bundle, timing) = reader
            .load_catalogue_for(&catalogue, role, self.generator_version)
            .await?;
        bundle.map(|bundle| (bundle, timing)).ok_or_else(|| {
            ApplicationError::NotFound(format!("node '{}' not found", request.root_node_id()))
        })
    }

    pub(crate) async fn materialize_selected_details(
        &self,
        bundle: KmpBundle,
        selected: &std::collections::BTreeSet<String>,
    ) -> Result<KmpBundle, ApplicationError> {
        let ids = std::iter::once(bundle.root_node())
            .chain(bundle.neighbor_nodes())
            .map(|node| node.node_id())
            .filter(|id| selected.contains(*id))
            .map(str::to_string)
            .collect::<Vec<_>>();
        let details = if ids.is_empty() {
            Vec::new()
        } else {
            self.detail_reader
                .load_node_details_batch(ids)
                .await?
                .into_iter()
                .flatten()
                .map(|detail| kmp_domain::BundleNodeDetail::from_projection(&detail))
                .collect()
        };
        Ok(bundle.with_node_details(details)?)
    }

    pub async fn get_context(
        &self,
        query: GetContextQuery,
    ) -> Result<GetContextResult, ApplicationError> {
        let rehydrate = RehydrateSessionUseCase::new(
            std::sync::Arc::clone(&self.graph_reader),
            std::sync::Arc::clone(&self.detail_reader),
            std::sync::Arc::clone(&self.snapshot_store),
            self.generator_version,
        );

        GetContextUseCase::new(rehydrate)
            .execute(
                &query.root_node_id,
                &query.role,
                query.depth,
                &query.requested_scopes,
                &query.render_options,
            )
            .await
    }
}