Skip to main content

kmp_application/queries/
get_context.rs

1use kmp_domain::{
2    GraphNeighborhoodReader, KmpBundle, NeighborhoodRequest, NodeDetailReader, SnapshotSaveOptions,
3    SnapshotStore,
4};
5
6use crate::ApplicationError;
7pub use crate::queries::render_graph_bundle::RenderedContext;
8use crate::queries::{
9    ContextRenderOptions, QueryApplicationService, QueryTimingBreakdown, RehydrateSessionUseCase,
10    clamp_native_graph_traversal_depth, render_graph_bundle_with_options,
11};
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct GetContextQuery {
15    pub root_node_id: String,
16    pub role: String,
17    pub depth: u32,
18    pub requested_scopes: Vec<String>,
19    pub render_options: ContextRenderOptions,
20}
21
22#[derive(Debug, Clone, PartialEq)]
23pub struct GetContextResult {
24    pub bundle: KmpBundle,
25    /// Identity of the consistent operation read, when the backend supplies it.
26    /// Derived readers must bypass revision caches when this is absent.
27    pub read_revision: Option<kmp_domain::GraphReadRevision>,
28    pub rendered: RenderedContext,
29    pub requested_scopes: Vec<String>,
30    pub served_at: std::time::SystemTime,
31    pub timing: Option<QueryTimingBreakdown>,
32}
33
34#[derive(Debug)]
35pub struct GetContextUseCase<G, D, S> {
36    rehydrate_session: RehydrateSessionUseCase<G, D, S>,
37}
38
39impl<G, D, S> GetContextUseCase<G, D, S>
40where
41    G: GraphNeighborhoodReader + Send + Sync,
42    D: NodeDetailReader + Send + Sync,
43    S: SnapshotStore + Send + Sync,
44{
45    pub fn new(rehydrate_session: RehydrateSessionUseCase<G, D, S>) -> Self {
46        Self { rehydrate_session }
47    }
48
49    pub async fn execute(
50        &self,
51        root_node_id: &str,
52        role: &str,
53        depth: u32,
54        requested_scopes: &[String],
55        render_options: &ContextRenderOptions,
56    ) -> Result<GetContextResult, ApplicationError> {
57        // The scopes were resolved before this query was built. Handing them to
58        // the reader is what lets a store narrow on the axis rather than the
59        // caller discarding what it should never have loaded.
60        let (bundle, timing) = self
61            .rehydrate_session
62            .execute_for(
63                &NeighborhoodRequest::new(root_node_id, clamp_native_graph_traversal_depth(depth))
64                    .with_scopes(requested_scopes.to_vec()),
65                role,
66                false,
67                SnapshotSaveOptions::default(),
68            )
69            .await?;
70        let rendered = render_graph_bundle_with_options(&bundle, render_options);
71
72        Ok(GetContextResult {
73            bundle,
74            read_revision: None,
75            rendered,
76            requested_scopes: requested_scopes.to_vec(),
77            served_at: std::time::SystemTime::now(),
78            timing: Some(timing),
79        })
80    }
81}
82
83impl<G, D, S> QueryApplicationService<G, D, S>
84where
85    G: GraphNeighborhoodReader + Send + Sync,
86    D: NodeDetailReader + Send + Sync,
87    S: SnapshotStore + Send + Sync,
88{
89    /// Read the structured context for consumers that do their own projection.
90    /// No prompt rendering, tokenization or snapshot write is needed here.
91    pub(crate) async fn read_context_bundle(
92        &self,
93        request: &NeighborhoodRequest,
94        role: &str,
95    ) -> Result<KmpBundle, ApplicationError> {
96        let rehydrate = RehydrateSessionUseCase::new(
97            std::sync::Arc::clone(&self.graph_reader),
98            std::sync::Arc::clone(&self.detail_reader),
99            std::sync::Arc::clone(&self.snapshot_store),
100            self.generator_version,
101        );
102        let (bundle, _) = rehydrate
103            .execute_for(request, role, false, SnapshotSaveOptions::default())
104            .await?;
105        Ok(bundle)
106    }
107
108    /// Read graph metadata first; the caller selects canonical bodies within
109    /// this same operation snapshot before returning a materialized result.
110    pub(crate) async fn read_context_catalogue(
111        &self,
112        request: &NeighborhoodRequest,
113        role: &str,
114    ) -> Result<KmpBundle, ApplicationError> {
115        Ok(self
116            .read_context_catalogue_with_timing(request, role)
117            .await?
118            .0)
119    }
120
121    pub(crate) async fn read_context_catalogue_with_timing(
122        &self,
123        request: &NeighborhoodRequest,
124        role: &str,
125    ) -> Result<(KmpBundle, QueryTimingBreakdown), ApplicationError> {
126        let reader = crate::queries::NodeCentricProjectionReader::new(
127            std::sync::Arc::clone(&self.graph_reader),
128            std::sync::Arc::clone(&self.detail_reader),
129        );
130        // Lane hints cannot remove other coordinates before whole-entry
131        // predicates run. Root selection already restricts the about scope.
132        // The full reader also retains this catalogue; defer bodies only.
133        let catalogue = NeighborhoodRequest::new(request.root_node_id(), request.depth());
134        let (bundle, timing) = reader
135            .load_catalogue_for(&catalogue, role, self.generator_version)
136            .await?;
137        bundle.map(|bundle| (bundle, timing)).ok_or_else(|| {
138            ApplicationError::NotFound(format!("node '{}' not found", request.root_node_id()))
139        })
140    }
141
142    pub(crate) async fn materialize_selected_details(
143        &self,
144        bundle: KmpBundle,
145        selected: &std::collections::BTreeSet<String>,
146    ) -> Result<KmpBundle, ApplicationError> {
147        let ids = std::iter::once(bundle.root_node())
148            .chain(bundle.neighbor_nodes())
149            .map(|node| node.node_id())
150            .filter(|id| selected.contains(*id))
151            .map(str::to_string)
152            .collect::<Vec<_>>();
153        let details = if ids.is_empty() {
154            Vec::new()
155        } else {
156            self.detail_reader
157                .load_node_details_batch(ids)
158                .await?
159                .into_iter()
160                .flatten()
161                .map(|detail| kmp_domain::BundleNodeDetail::from_projection(&detail))
162                .collect()
163        };
164        Ok(bundle.with_node_details(details)?)
165    }
166
167    pub async fn get_context(
168        &self,
169        query: GetContextQuery,
170    ) -> Result<GetContextResult, ApplicationError> {
171        let rehydrate = RehydrateSessionUseCase::new(
172            std::sync::Arc::clone(&self.graph_reader),
173            std::sync::Arc::clone(&self.detail_reader),
174            std::sync::Arc::clone(&self.snapshot_store),
175            self.generator_version,
176        );
177
178        GetContextUseCase::new(rehydrate)
179            .execute(
180                &query.root_node_id,
181                &query.role,
182                query.depth,
183                &query.requested_scopes,
184                &query.render_options,
185            )
186            .await
187    }
188}