1use crate::{ContextQuery, ContextResponse, MCPToolResult, get_memory_system, parse_datetime};
4use anyhow::Result;
5use post_cortex_core::core::context_update::EntityType;
6use post_cortex_core::core::context_update::UpdateType;
7use post_cortex_core::core::timeout_utils::with_mcp_timeout;
8use post_cortex_core::session::active_session::ActiveSession;
9use post_cortex_memory::ConversationMemorySystem;
10use std::collections::HashMap;
11use tracing::{debug, error};
12use uuid::Uuid;
13
14pub async fn query_conversation_context_with_system(
16 query_type: String,
17 parameters: HashMap<String, String>,
18 session_id: Uuid,
19 system: &ConversationMemorySystem,
20) -> Result<MCPToolResult> {
21 eprintln!(
22 "DEBUG: query_conversation_context_with_system - Looking for session: {}",
23 session_id
24 );
25
26 let result = with_mcp_timeout(async {
27 let session_arc = match system.get_session(session_id).await {
28 Ok(session) => {
29 eprintln!(
30 "DEBUG: query_conversation_context_with_system - Session found successfully"
31 );
32 session
33 }
34 Err(e) => {
35 eprintln!(
36 "DEBUG: query_conversation_context_with_system - Session not found: {}",
37 e
38 );
39 return Err(anyhow::anyhow!("Session not found: {}", e));
40 }
41 };
42 let session = session_arc.load();
43 debug!("query_conversation_context: session {} loaded", session_id);
44
45 let query = match query_type.as_str() {
46 "recent_changes" => {
47 let since_str = parameters.get("since").cloned().unwrap_or_default();
48 let since = parse_datetime(&since_str)?;
49 ContextQuery::GetRecentChanges { since }
50 }
51 "code_references" => {
52 let file_path = parameters.get("file_path").cloned().unwrap_or_default();
53 ContextQuery::FindCodeReferences { file_path }
54 }
55 "structured_summary" => ContextQuery::GetStructuredSummary,
56 "decisions" => ContextQuery::GetDecisions { since: None },
57 "open_questions" => ContextQuery::GetOpenQuestions,
58 "related_entities" => {
59 let entity_name = parameters.get("entity_name").cloned().unwrap_or_default();
60 ContextQuery::FindRelatedEntities { entity_name }
61 }
62 "entity_context" => {
63 let entity_name = parameters.get("entity_name").cloned().unwrap_or_default();
64 ContextQuery::GetEntityContext { entity_name }
65 }
66 "all_entities" => ContextQuery::GetAllEntities { entity_type: None },
67 "trace_relationships" => {
68 let from_entity = parameters.get("entity_name").cloned().unwrap_or_default();
69 let max_depth = parameters
70 .get("max_depth")
71 .and_then(|s| s.parse().ok())
72 .unwrap_or(3);
73 ContextQuery::TraceRelationships {
74 from_entity,
75 max_depth,
76 }
77 }
78 "find_related_entities" => {
79 let entity_name = parameters.get("entity_name").cloned().unwrap_or_default();
80 ContextQuery::FindRelatedEntities { entity_name }
81 }
82 "get_entity_context" => {
83 let entity_name = parameters.get("entity_name").cloned().unwrap_or_default();
84 ContextQuery::GetEntityContext { entity_name }
85 }
86 "get_entity_network" => {
87 let center_entity = parameters.get("entity_name").cloned().unwrap_or_default();
88 let max_depth = parameters
89 .get("max_depth")
90 .and_then(|s| s.parse().ok())
91 .unwrap_or(2);
92 ContextQuery::GetEntityNetwork {
93 center_entity,
94 max_depth,
95 }
96 }
97 "get_most_important_entities" => {
98 let limit = parameters
99 .get("limit")
100 .and_then(|s| s.parse().ok())
101 .unwrap_or(10);
102 ContextQuery::GetMostImportantEntities { limit }
103 }
104 "get_recently_mentioned_entities" => {
105 let limit = parameters
106 .get("limit")
107 .and_then(|s| s.parse().ok())
108 .unwrap_or(10);
109 ContextQuery::GetRecentlyMentionedEntities { limit }
110 }
111 "analyze_entity_importance" => ContextQuery::AnalyzeEntityImportance,
112 "find_entities_by_type" => {
113 let entity_type_str = parameters.get("entity_type").cloned().unwrap_or_default();
114 let entity_type = match entity_type_str.as_str() {
115 "technology" => EntityType::Technology,
116 "concept" => EntityType::Concept,
117 "problem" => EntityType::Problem,
118 "solution" => EntityType::Solution,
119 "decision" => EntityType::Decision,
120 "code_component" => EntityType::CodeComponent,
121 _ => EntityType::Concept,
122 };
123 ContextQuery::FindEntitiesByType { entity_type }
124 }
125 "search_updates" => {
126 let query = parameters.get("query").cloned().unwrap_or_default();
127 ContextQuery::SearchUpdates { query }
128 }
129 "assemble_context" => {
130 let query = parameters.get("query").cloned().unwrap_or_default();
131 let token_budget = parameters
132 .get("token_budget")
133 .and_then(|s| s.parse().ok())
134 .unwrap_or(4000);
135 ContextQuery::AssembleContext {
136 query,
137 token_budget,
138 }
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
173pub 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 {
277 query,
278 token_budget,
279 }
280 }
281 _ => {
282 return Ok(MCPToolResult::error(format!(
283 "Unknown query type: {}",
284 query_type
285 )));
286 }
287 };
288
289 let response = query_context(&session, query).await?;
290 let json_response = serde_json::to_value(response)?;
291
292 Ok(MCPToolResult::success(
293 "Query successful".to_string(),
294 Some(json_response),
295 ))
296}
297
298pub(crate) async fn query_context(
300 session: &ActiveSession,
301 query: ContextQuery,
302) -> Result<ContextResponse> {
303 use post_cortex_core::core::context_update::CodeReference;
304
305 match query {
306 ContextQuery::GetRecentChanges { since } => {
307 let recent_updates: Vec<post_cortex_core::core::context_update::ContextUpdate> =
308 session
309 .hot_context
310 .iter()
311 .iter()
312 .chain(session.warm_context.iter().map(|c| &c.update))
313 .filter(|u| u.timestamp >= since)
314 .cloned()
315 .collect();
316 Ok(ContextResponse::RecentChanges(recent_updates))
317 }
318 ContextQuery::FindCodeReferences { file_path } => {
319 let refs = session
320 .code_references
321 .get(&file_path)
322 .cloned()
323 .unwrap_or_default();
324 let converted_refs: Vec<CodeReference> = refs
325 .into_iter()
326 .map(|r| CodeReference {
327 file_path: r.file_path,
328 start_line: r.start_line,
329 end_line: r.end_line,
330 code_snippet: r.code_snippet,
331 commit_hash: r.commit_hash,
332 branch: r.branch,
333 change_description: r.change_description,
334 })
335 .collect();
336 Ok(ContextResponse::CodeReferences(converted_refs))
337 }
338 ContextQuery::GetStructuredSummary => Ok(ContextResponse::StructuredSummary(
339 (*session.current_state).clone(),
340 )),
341 ContextQuery::FindRelatedEntities { entity_name } => {
342 let related = session.entity_graph.find_related_entities(&entity_name);
343 Ok(ContextResponse::RelatedEntities(related))
344 }
345 ContextQuery::GetEntityContext { entity_name } => {
346 let context = session.entity_graph.get_entity_context(&entity_name);
347 Ok(ContextResponse::EntityContext(
348 context.unwrap_or("Entity not found".to_string()),
349 ))
350 }
351 ContextQuery::GetAllEntities { entity_type } => {
352 let entities: Vec<String> = match entity_type {
353 Some(et) => session
354 .entity_graph
355 .get_entities_by_type(&et)
356 .into_iter()
357 .map(|e| e.name.clone())
358 .collect(),
359 None => session
360 .entity_graph
361 .get_most_important_entities(50)
362 .into_iter()
363 .map(|e| e.name.clone())
364 .collect(),
365 };
366 Ok(ContextResponse::AllEntities(entities))
367 }
368 ContextQuery::TraceRelationships {
369 from_entity,
370 max_depth,
371 } => {
372 let trace = session
373 .entity_graph
374 .trace_entity_relationships(&from_entity, max_depth);
375 Ok(ContextResponse::Entities(
376 trace.into_iter().map(|(a, _, _)| a).collect(),
377 ))
378 }
379 ContextQuery::GetEntityNetwork {
380 center_entity,
381 max_depth,
382 } => {
383 let _network = session
384 .entity_graph
385 .get_entity_network(¢er_entity, max_depth);
386 Ok(ContextResponse::EntityNetwork("Network data".to_string()))
387 }
388 ContextQuery::GetMostImportantEntities { limit } => {
389 let entities = session.entity_graph.get_most_important_entities(limit);
390 Ok(ContextResponse::Entities(
391 entities.into_iter().map(|e| e.name.clone()).collect(),
392 ))
393 }
394 ContextQuery::GetRecentlyMentionedEntities { limit } => {
395 let entities = session.entity_graph.get_recently_mentioned_entities(limit);
396 Ok(ContextResponse::Entities(
397 entities.into_iter().map(|e| e.name.clone()).collect(),
398 ))
399 }
400 ContextQuery::AnalyzeEntityImportance => {
401 let _analysis = session.entity_graph.analyze_entity_importance();
402 Ok(ContextResponse::ImportanceAnalysis(
403 "Analysis complete".to_string(),
404 ))
405 }
406 ContextQuery::FindEntitiesByType { entity_type } => {
407 let entities = session.entity_graph.get_entities_by_type(&entity_type);
408 Ok(ContextResponse::Entities(
409 entities.into_iter().map(|e| e.name.clone()).collect(),
410 ))
411 }
412 ContextQuery::SearchUpdates { query } => {
413 let update_results: Vec<post_cortex_core::core::context_update::ContextUpdate> =
414 session
415 .hot_context
416 .iter()
417 .iter()
418 .chain(session.warm_context.iter().map(|c| &c.update))
419 .filter(|u| {
420 u.content
421 .title
422 .to_lowercase()
423 .contains(&query.to_lowercase())
424 || u.content
425 .description
426 .to_lowercase()
427 .contains(&query.to_lowercase())
428 })
429 .cloned()
430 .collect();
431 Ok(ContextResponse::SearchResults(update_results))
432 }
433 ContextQuery::GetDecisions { since: _ } => {
434 let decisions: Vec<post_cortex_core::core::context_update::ContextUpdate> = session
435 .hot_context
436 .iter()
437 .into_iter()
438 .filter(|u| matches!(u.update_type, UpdateType::DecisionMade))
439 .collect();
440 Ok(ContextResponse::Decisions(decisions))
441 }
442 ContextQuery::GetOpenQuestions => Ok(ContextResponse::OpenQuestions(vec![
443 "No open questions".to_string(),
444 ])),
445 ContextQuery::GetChangeHistory { file_path: _ } => {
446 let changes: Vec<post_cortex_core::core::context_update::ContextUpdate> = session
447 .hot_context
448 .iter()
449 .into_iter()
450 .filter(|u| matches!(u.update_type, UpdateType::CodeChanged))
451 .collect();
452 Ok(ContextResponse::ChangeHistory(changes))
453 }
454 ContextQuery::AssembleContext {
455 query,
456 token_budget,
457 } => {
458 use post_cortex_memory::context_assembly;
459
460 let updates: Vec<_> = session
461 .hot_context
462 .iter()
463 .iter()
464 .chain(session.warm_context.iter().map(|c| &c.update))
465 .cloned()
466 .collect();
467
468 let assembled = context_assembly::assemble_context(
469 &query,
470 &session.entity_graph,
471 &updates,
472 token_budget,
473 );
474
475 Ok(ContextResponse::AssembledContext(assembled))
476 }
477 _ => Ok(ContextResponse::Entities(vec![
478 "Not implemented".to_string(),
479 ])),
480 }
481}