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

kmp_application/queries/
get_context.rs
