post-cortex-daemon 0.5.0

// Copyright (c) 2025, 2026 Julius ML
// MIT License

//! Tool 7: assemble_context — graph-aware retrieval: query-term matching over
//! recent updates, entity-graph traversal and impact analysis.

use crate::daemon::assembly_scope;
use crate::daemon::coerce::{CoercionError, coerce_and_validate};
use crate::daemon::validate::validate_session_id;
use post_cortex_mcp::MCPToolResult;
use post_cortex_memory::ConversationMemorySystem;
use post_cortex_memory::context_assembly;
use rmcp::{
    handler::server::wrapper::Parameters,
    model::{CallToolResult, ErrorData as McpError},
};
use std::sync::Arc;
use uuid::Uuid;

use super::AssembleContextRequest;
use super::mcp_result_to_call_result;

pub(super) async fn handle(
    memory_system: &Arc<ConversationMemorySystem>,
    params: Parameters<serde_json::Value>,
) -> Result<CallToolResult, McpError> {
    let req: AssembleContextRequest = coerce_and_validate(params.0).map_err(|e| {
        if e.message.contains("query") {
            e.clone()
                .with_parameter_path("query".to_string())
                .with_expected_type("non-empty query/topic string")
                .with_hint("Provide the question or topic to assemble context for")
                .to_mcp_error()
        } else {
            e.to_mcp_error()
        }
    })?;

    if req.query.trim().is_empty() {
        return Err(McpError::invalid_params(
            "query cannot be empty".to_string(),
            Some(serde_json::Value::String("query".to_string())),
        ));
    }

    if req.session_id.is_none() && req.workspace_id.is_none() {
        return Err(McpError::invalid_params(
            "either session_id or workspace_id is required".to_string(),
            None,
        ));
    }

    let token_budget = req
        .token_budget
        .filter(|b| *b > 0)
        .map(|b| b as usize)
        .unwrap_or(4000);
    let corpus = post_cortex_mcp::search::parse_sources(req.sources.clone()).map_err(|e| {
        McpError::invalid_params(e, Some(serde_json::json!({"parameter": "sources"})))
    })?;

    let (updates, graph, sessions, workspaces) =
        if let Some(ws_raw) = req.workspace_id.as_ref().filter(|s| !s.is_empty()) {
            let ws_uuid = Uuid::parse_str(ws_raw).map_err(|_| {
                CoercionError::new(
                    "Invalid UUID",
                    std::io::Error::new(std::io::ErrorKind::InvalidInput, "Invalid UUID format"),
                    Some(serde_json::Value::String(ws_raw.clone())),
                )
                .with_parameter_path("workspace_id".to_string())
                .with_expected_type("Valid UUID string")
                .to_mcp_error()
            })?;

            let workspace = memory_system
                .workspace_manager
                .get_workspace(&ws_uuid)
                .ok_or_else(|| {
                    McpError::invalid_params(
                        format!("Workspace {} not found", ws_uuid),
                        Some(serde_json::Value::String("workspace_id".to_string())),
                    )
                })?;

            let mut all_updates = Vec::new();
            let mut graphs = Vec::new();
            let mut sessions = Vec::new();

            for (ws_session_id, _role) in workspace.get_all_sessions() {
                if let Ok(session_arc) = memory_system.get_session(ws_session_id).await {
                    let session = session_arc.load();
                    assembly_scope::session_updates(&session, &mut all_updates);
                    graphs.push((ws_session_id, Arc::clone(&session.entity_graph)));
                    sessions.push(ws_session_id);
                }
            }

            let merged = assembly_scope::WORKSPACE_GRAPHS.merged(ws_uuid, graphs);
            (all_updates, merged, sessions, vec![ws_uuid])
        } else {
            let sid_raw = req.session_id.as_ref().unwrap();
            let session_id = validate_session_id(sid_raw).map_err(|e| e.to_mcp_error())?;

            let session_arc = memory_system
                .get_session(session_id)
                .await
                .map_err(|e| McpError::internal_error(format!("Session not found: {}", e), None))?;
            let session = session_arc.load();
            let mut updates = Vec::new();
            assembly_scope::session_updates(&session, &mut updates);
            (
                updates,
                Arc::clone(&session.entity_graph),
                vec![session_id],
                assembly_scope::workspaces_of(memory_system, session_id),
            )
        };

    let assembled = assembly_scope::assemble(
        memory_system,
        &req.query,
        &graph,
        &updates,
        &sessions,
        &workspaces,
        token_budget,
        corpus,
    )
    .await;
    let formatted_text = context_assembly::format_for_llm(&assembled);

    let payload = serde_json::json!({
        "query": req.query,
        "token_budget": token_budget,
        "total_tokens": assembled.total_tokens,
        "items": assembled.items,
        "entity_context": assembled.entity_context,
        "impact": assembled.impact,
        "cards": assembled.cards,
        "formatted_text": formatted_text,
    });

    // Apply optional JSONPath select filter
    let payload = if req.select.is_some() || req.select_map.is_some() {
        use post_cortex_memory::jsonpath;

        let select_string = req.select.as_deref();
        let select_map = req.select_map.as_ref();

        match jsonpath::apply_select(&payload, select_string, select_map) {
            Ok(filtered) => filtered,
            Err(e) => {
                return Err(McpError::invalid_params(
                    format!("Invalid JSONPath expression: {}", e.message),
                    Some(serde_json::json!({
                        "offset": e.byte_offset,
                        "hint": "Allowed: $, .field, [n], [-n], [a:b], [*], [0,3,7]. \
                                 Forbidden: .., [?(...)], functions, [a:b:c]"
                    })),
                ));
            }
        }
    } else {
        payload
    };

    Ok(mcp_result_to_call_result(MCPToolResult::success(
        format!(
            "Assembled {} items ({} tokens) for query: {}",
            assembled.items.len(),
            assembled.total_tokens,
            req.query
        ),
        Some(payload),
    )))
}