post_cortex_mcp/
query.rs

1//! Structured and keyword-based queries over session context.
2
3use post_cortex_core::core::context_update::EntityType;
4use post_cortex_memory::ConversationMemorySystem;
5use post_cortex_core::core::timeout_utils::with_mcp_timeout;
6use post_cortex_core::session::active_session::ActiveSession;
7use crate::{
8    get_memory_system, parse_datetime, ContextQuery, ContextResponse,
9    MCPToolResult,
10};
11use post_cortex_core::core::context_update::UpdateType;
12use anyhow::Result;
13use std::collections::HashMap;
14use tracing::{debug, error};
15use uuid::Uuid;
16
17/// Execute a typed query against a session using an explicit memory system reference.
18pub async fn query_conversation_context_with_system(
19    query_type: String,
20    parameters: HashMap<String, String>,
21    session_id: Uuid,
22    system: &ConversationMemorySystem,
23) -> Result<MCPToolResult> {
24    eprintln!(
25        "DEBUG: query_conversation_context_with_system - Looking for session: {}",
26        session_id
27    );
28
29    let result = with_mcp_timeout(async {
30        let session_arc = match system.get_session(session_id).await {
31            Ok(session) => {
32                eprintln!(
33                    "DEBUG: query_conversation_context_with_system - Session found successfully"
34                );
35                session
36            }
37            Err(e) => {
38                eprintln!(
39                    "DEBUG: query_conversation_context_with_system - Session not found: {}",
40                    e
41                );
42                return Err(anyhow::anyhow!("Session not found: {}", e));
43            }
44        };
45        let session = session_arc.load();
46        debug!("query_conversation_context: session {} loaded", session_id);
47
48        let query = match query_type.as_str() {
49            "recent_changes" => {
50                let since_str = parameters.get("since").cloned().unwrap_or_default();
51                let since = parse_datetime(&since_str)?;
52                ContextQuery::GetRecentChanges { since }
53            }
54            "code_references" => {
55                let file_path = parameters.get("file_path").cloned().unwrap_or_default();
56                ContextQuery::FindCodeReferences { file_path }
57            }
58            "structured_summary" => ContextQuery::GetStructuredSummary,
59            "decisions" => ContextQuery::GetDecisions { since: None },
60            "open_questions" => ContextQuery::GetOpenQuestions,
61            "related_entities" => {
62                let entity_name = parameters.get("entity_name").cloned().unwrap_or_default();
63                ContextQuery::FindRelatedEntities { entity_name }
64            }
65            "entity_context" => {
66                let entity_name = parameters.get("entity_name").cloned().unwrap_or_default();
67                ContextQuery::GetEntityContext { entity_name }
68            }
69            "all_entities" => ContextQuery::GetAllEntities { entity_type: None },
70            "trace_relationships" => {
71                let from_entity = parameters.get("entity_name").cloned().unwrap_or_default();
72                let max_depth = parameters
73                    .get("max_depth")
74                    .and_then(|s| s.parse().ok())
75                    .unwrap_or(3);
76                ContextQuery::TraceRelationships {
77                    from_entity,
78                    max_depth,
79                }
80            }
81            "find_related_entities" => {
82                let entity_name = parameters.get("entity_name").cloned().unwrap_or_default();
83                ContextQuery::FindRelatedEntities { entity_name }
84            }
85            "get_entity_context" => {
86                let entity_name = parameters.get("entity_name").cloned().unwrap_or_default();
87                ContextQuery::GetEntityContext { entity_name }
88            }
89            "get_entity_network" => {
90                let center_entity = parameters.get("entity_name").cloned().unwrap_or_default();
91                let max_depth = parameters
92                    .get("max_depth")
93                    .and_then(|s| s.parse().ok())
94                    .unwrap_or(2);
95                ContextQuery::GetEntityNetwork {
96                    center_entity,
97                    max_depth,
98                }
99            }
100            "get_most_important_entities" => {
101                let limit = parameters
102                    .get("limit")
103                    .and_then(|s| s.parse().ok())
104                    .unwrap_or(10);
105                ContextQuery::GetMostImportantEntities { limit }
106            }
107            "get_recently_mentioned_entities" => {
108                let limit = parameters
109                    .get("limit")
110                    .and_then(|s| s.parse().ok())
111                    .unwrap_or(10);
112                ContextQuery::GetRecentlyMentionedEntities { limit }
113            }
114            "analyze_entity_importance" => ContextQuery::AnalyzeEntityImportance,
115            "find_entities_by_type" => {
116                let entity_type_str = parameters.get("entity_type").cloned().unwrap_or_default();
117                let entity_type = match entity_type_str.as_str() {
118                    "technology" => EntityType::Technology,
119                    "concept" => EntityType::Concept,
120                    "problem" => EntityType::Problem,
121                    "solution" => EntityType::Solution,
122                    "decision" => EntityType::Decision,
123                    "code_component" => EntityType::CodeComponent,
124                    _ => EntityType::Concept,
125                };
126                ContextQuery::FindEntitiesByType { entity_type }
127            }
128            "search_updates" => {
129                let query = parameters.get("query").cloned().unwrap_or_default();
130                ContextQuery::SearchUpdates { query }
131            }
132            "assemble_context" => {
133                let query = parameters.get("query").cloned().unwrap_or_default();
134                let token_budget = parameters
135                    .get("token_budget")
136                    .and_then(|s| s.parse().ok())
137                    .unwrap_or(4000);
138                ContextQuery::AssembleContext { query, token_budget }
139            }
140            _ => {
141                return Ok(MCPToolResult::error(format!(
142                    "Unknown query type: {}",
143                    query_type
144                )));
145            }
146        };
147
148        let response = query_context(&session, query).await?;
149        let json_response = serde_json::to_value(response)?;
150
151        Ok(MCPToolResult::success(
152            "Query successful".to_string(),
153            Some(json_response),
154        ))
155    })
156    .await;
157
158    match result {
159        Ok(success_result) => success_result,
160        Err(timeout_error) => {
161            error!(
162                "TIMEOUT: query_conversation_context_with_system - session: {}, error: {}",
163                session_id, timeout_error
164            );
165            Ok(MCPToolResult::error(format!(
166                "Query timed out: {}",
167                timeout_error
168            )))
169        }
170    }
171}
172
173/// Execute a typed query against a session via the global memory system.
174pub async fn query_conversation_context(
175    query_type: String,
176    parameters: HashMap<String, String>,
177    session_id: Uuid,
178) -> Result<MCPToolResult> {
179    let system = get_memory_system().await?;
180    let session_arc = system
181        .get_session(session_id)
182        .await
183        .map_err(|e| anyhow::anyhow!("Failed to load session: {}", e))?;
184    let session = session_arc.load();
185
186    let query = match query_type.as_str() {
187        "recent_changes" => {
188            let since_str = parameters.get("since").cloned().unwrap_or_default();
189            let since = parse_datetime(&since_str)?;
190            ContextQuery::GetRecentChanges { since }
191        }
192        "code_references" => {
193            let file_path = parameters.get("file_path").cloned().unwrap_or_default();
194            ContextQuery::FindCodeReferences { file_path }
195        }
196        "structured_summary" => ContextQuery::GetStructuredSummary,
197        "decisions" => ContextQuery::GetDecisions { since: None },
198        "open_questions" => ContextQuery::GetOpenQuestions,
199        "related_entities" => {
200            let entity_name = parameters.get("entity_name").cloned().unwrap_or_default();
201            ContextQuery::FindRelatedEntities { entity_name }
202        }
203        "entity_context" => {
204            let entity_name = parameters.get("entity_name").cloned().unwrap_or_default();
205            ContextQuery::GetEntityContext { entity_name }
206        }
207        "all_entities" => ContextQuery::GetAllEntities { entity_type: None },
208        "trace_relationships" => {
209            let from_entity = parameters.get("entity_name").cloned().unwrap_or_default();
210            let max_depth = parameters
211                .get("max_depth")
212                .and_then(|s| s.parse().ok())
213                .unwrap_or(3);
214            ContextQuery::TraceRelationships {
215                from_entity,
216                max_depth,
217            }
218        }
219        "find_related_entities" => {
220            let entity_name = parameters.get("entity_name").cloned().unwrap_or_default();
221            ContextQuery::FindRelatedEntities { entity_name }
222        }
223        "get_entity_context" => {
224            let entity_name = parameters.get("entity_name").cloned().unwrap_or_default();
225            ContextQuery::GetEntityContext { entity_name }
226        }
227        "get_entity_network" => {
228            let center_entity = parameters.get("entity_name").cloned().unwrap_or_default();
229            let max_depth = parameters
230                .get("max_depth")
231                .and_then(|s| s.parse().ok())
232                .unwrap_or(2);
233            ContextQuery::GetEntityNetwork {
234                center_entity,
235                max_depth,
236            }
237        }
238        "get_most_important_entities" => {
239            let limit = parameters
240                .get("limit")
241                .and_then(|s| s.parse().ok())
242                .unwrap_or(10);
243            ContextQuery::GetMostImportantEntities { limit }
244        }
245        "get_recently_mentioned_entities" => {
246            let limit = parameters
247                .get("limit")
248                .and_then(|s| s.parse().ok())
249                .unwrap_or(10);
250            ContextQuery::GetRecentlyMentionedEntities { limit }
251        }
252        "analyze_entity_importance" => ContextQuery::AnalyzeEntityImportance,
253        "find_entities_by_type" => {
254            let entity_type_str = parameters.get("entity_type").cloned().unwrap_or_default();
255            let entity_type = match entity_type_str.as_str() {
256                "technology" => EntityType::Technology,
257                "concept" => EntityType::Concept,
258                "problem" => EntityType::Problem,
259                "solution" => EntityType::Solution,
260                "decision" => EntityType::Decision,
261                "code_component" => EntityType::CodeComponent,
262                _ => EntityType::Concept,
263            };
264            ContextQuery::FindEntitiesByType { entity_type }
265        }
266        "search_updates" => {
267            let query = parameters.get("query").cloned().unwrap_or_default();
268            ContextQuery::SearchUpdates { query }
269        }
270        "assemble_context" => {
271            let query = parameters.get("query").cloned().unwrap_or_default();
272            let token_budget = parameters
273                .get("token_budget")
274                .and_then(|s| s.parse().ok())
275                .unwrap_or(4000);
276            ContextQuery::AssembleContext { query, token_budget }
277        }
278        _ => {
279            return Ok(MCPToolResult::error(format!(
280                "Unknown query type: {}",
281                query_type
282            )));
283        }
284    };
285
286    let response = query_context(&session, query).await?;
287    let json_response = serde_json::to_value(response)?;
288
289    Ok(MCPToolResult::success(
290        "Query successful".to_string(),
291        Some(json_response),
292    ))
293}
294
295/// Dispatch a [`ContextQuery`] against a loaded session and return a [`ContextResponse`].
296pub(crate) async fn query_context(
297    session: &ActiveSession,
298    query: ContextQuery,
299) -> Result<ContextResponse> {
300    use post_cortex_core::core::context_update::CodeReference;
301
302    match query {
303        ContextQuery::GetRecentChanges { since } => {
304            let recent_updates: Vec<post_cortex_core::core::context_update::ContextUpdate> = session
305                .hot_context
306                .iter()
307                .iter()
308                .chain(session.warm_context.iter().map(|c| &c.update))
309                .filter(|u| u.timestamp >= since)
310                .cloned()
311                .collect();
312            Ok(ContextResponse::RecentChanges(recent_updates))
313        }
314        ContextQuery::FindCodeReferences { file_path } => {
315            let refs = session
316                .code_references
317                .get(&file_path)
318                .cloned()
319                .unwrap_or_default();
320            let converted_refs: Vec<CodeReference> = refs
321                .into_iter()
322                .map(|r| CodeReference {
323                    file_path: r.file_path,
324                    start_line: r.start_line,
325                    end_line: r.end_line,
326                    code_snippet: r.code_snippet,
327                    commit_hash: r.commit_hash,
328                    branch: r.branch,
329                    change_description: r.change_description,
330                })
331                .collect();
332            Ok(ContextResponse::CodeReferences(converted_refs))
333        }
334        ContextQuery::GetStructuredSummary => Ok(ContextResponse::StructuredSummary(
335            (*session.current_state).clone(),
336        )),
337        ContextQuery::FindRelatedEntities { entity_name } => {
338            let related = session.entity_graph.find_related_entities(&entity_name);
339            Ok(ContextResponse::RelatedEntities(related))
340        }
341        ContextQuery::GetEntityContext { entity_name } => {
342            let context = session.entity_graph.get_entity_context(&entity_name);
343            Ok(ContextResponse::EntityContext(
344                context.unwrap_or("Entity not found".to_string()),
345            ))
346        }
347        ContextQuery::GetAllEntities { entity_type } => {
348            let entities: Vec<String> = match entity_type {
349                Some(et) => session
350                    .entity_graph
351                    .get_entities_by_type(&et)
352                    .into_iter()
353                    .map(|e| e.name.clone())
354                    .collect(),
355                None => session
356                    .entity_graph
357                    .get_most_important_entities(50)
358                    .into_iter()
359                    .map(|e| e.name.clone())
360                    .collect(),
361            };
362            Ok(ContextResponse::AllEntities(entities))
363        }
364        ContextQuery::TraceRelationships {
365            from_entity,
366            max_depth,
367        } => {
368            let trace = session
369                .entity_graph
370                .trace_entity_relationships(&from_entity, max_depth);
371            Ok(ContextResponse::Entities(
372                trace.into_iter().map(|(a, _, _)| a).collect(),
373            ))
374        }
375        ContextQuery::GetEntityNetwork {
376            center_entity,
377            max_depth,
378        } => {
379            let _network = session
380                .entity_graph
381                .get_entity_network(&center_entity, max_depth);
382            Ok(ContextResponse::EntityNetwork("Network data".to_string()))
383        }
384        ContextQuery::GetMostImportantEntities { limit } => {
385            let entities = session.entity_graph.get_most_important_entities(limit);
386            Ok(ContextResponse::Entities(
387                entities.into_iter().map(|e| e.name.clone()).collect(),
388            ))
389        }
390        ContextQuery::GetRecentlyMentionedEntities { limit } => {
391            let entities = session.entity_graph.get_recently_mentioned_entities(limit);
392            Ok(ContextResponse::Entities(
393                entities.into_iter().map(|e| e.name.clone()).collect(),
394            ))
395        }
396        ContextQuery::AnalyzeEntityImportance => {
397            let _analysis = session.entity_graph.analyze_entity_importance();
398            Ok(ContextResponse::ImportanceAnalysis(
399                "Analysis complete".to_string(),
400            ))
401        }
402        ContextQuery::FindEntitiesByType { entity_type } => {
403            let entities = session.entity_graph.get_entities_by_type(&entity_type);
404            Ok(ContextResponse::Entities(
405                entities.into_iter().map(|e| e.name.clone()).collect(),
406            ))
407        }
408        ContextQuery::SearchUpdates { query } => {
409            let update_results: Vec<post_cortex_core::core::context_update::ContextUpdate> = session
410                .hot_context
411                .iter()
412                .iter()
413                .chain(session.warm_context.iter().map(|c| &c.update))
414                .filter(|u| {
415                    u.content
416                        .title
417                        .to_lowercase()
418                        .contains(&query.to_lowercase())
419                        || u.content
420                            .description
421                            .to_lowercase()
422                            .contains(&query.to_lowercase())
423                })
424                .cloned()
425                .collect();
426            Ok(ContextResponse::SearchResults(update_results))
427        }
428        ContextQuery::GetDecisions { since: _ } => {
429            let decisions: Vec<post_cortex_core::core::context_update::ContextUpdate> = session
430                .hot_context
431                .iter()
432                .into_iter()
433                .filter(|u| matches!(u.update_type, UpdateType::DecisionMade))
434                .collect();
435            Ok(ContextResponse::Decisions(decisions))
436        }
437        ContextQuery::GetOpenQuestions => Ok(ContextResponse::OpenQuestions(vec![
438            "No open questions".to_string(),
439        ])),
440        ContextQuery::GetChangeHistory { file_path: _ } => {
441            let changes: Vec<post_cortex_core::core::context_update::ContextUpdate> = session
442                .hot_context
443                .iter()
444                .into_iter()
445                .filter(|u| matches!(u.update_type, UpdateType::CodeChanged))
446                .collect();
447            Ok(ContextResponse::ChangeHistory(changes))
448        }
449        ContextQuery::AssembleContext { query, token_budget } => {
450            use post_cortex_memory::context_assembly;
451
452            let updates: Vec<_> = session.hot_context.iter().iter()
453                .chain(session.warm_context.iter().map(|c| &c.update))
454                .cloned()
455                .collect();
456
457            let assembled = context_assembly::assemble_context(
458                &query,
459                &session.entity_graph,
460                &updates,
461                token_budget,
462            );
463
464            Ok(ContextResponse::AssembledContext(assembled))
465        }
466        _ => Ok(ContextResponse::Entities(vec![
467            "Not implemented".to_string(),
468        ])),
469    }
470}
post_cortex_mcp/query.rs

post_cortex_mcp/
query.rs
