backend/app/mcp_server.py

"""MCP Server for Project Memory - The Brain of the system."""

import asyncio
import json
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent

from app.database import init_db
from app.tools.projects import create_project, list_projects, join_project
from app.tools.tasks import create_task, list_tasks, list_activity
from app.tools.memory import complete_task, memory_search
from app.smart_query import smart_query
from app.vectorstore import init_vectorstore

# Initialize the MCP server
server = Server("project-memory")


# Tool definitions
TOOLS = [
    # Project Tools
    Tool(
        name="create_project",
        description="Create a new project and add the creator as owner",
        inputSchema={
            "type": "object",
            "properties": {
                "name": {"type": "string", "description": "Project name"},
                "description": {"type": "string", "description": "Project description"},
                "userId": {"type": "string", "description": "ID of the user creating the project"}
            },
            "required": ["name", "description", "userId"]
        }
    ),
    Tool(
        name="list_projects",
        description="List all projects for a user",
        inputSchema={
            "type": "object",
            "properties": {
                "userId": {"type": "string", "description": "ID of the user"}
            },
            "required": ["userId"]
        }
    ),
    Tool(
        name="join_project",
        description="Add a user to an existing project",
        inputSchema={
            "type": "object",
            "properties": {
                "projectId": {"type": "string", "description": "ID of the project to join"},
                "userId": {"type": "string", "description": "ID of the user joining"}
            },
            "required": ["projectId", "userId"]
        }
    ),
    # Task Tools
    Tool(
        name="create_task",
        description="Create a new task in a project",
        inputSchema={
            "type": "object",
            "properties": {
                "projectId": {"type": "string", "description": "ID of the project"},
                "title": {"type": "string", "description": "Task title"},
                "description": {"type": "string", "description": "Task description"},
                "assignedTo": {"type": "string", "description": "User ID to assign the task to"}
            },
            "required": ["projectId", "title"]
        }
    ),
    Tool(
        name="list_tasks",
        description="List all tasks for a project, optionally filtered by status",
        inputSchema={
            "type": "object",
            "properties": {
                "projectId": {"type": "string", "description": "ID of the project"},
                "status": {
                    "type": "string",
                    "enum": ["todo", "in_progress", "done"],
                    "description": "Optional status filter"
                }
            },
            "required": ["projectId"]
        }
    ),
    Tool(
        name="list_activity",
        description="Get recent activity (log entries) for a project",
        inputSchema={
            "type": "object",
            "properties": {
                "projectId": {"type": "string", "description": "ID of the project"},
                "limit": {"type": "number", "default": 20, "description": "Maximum number of entries"}
            },
            "required": ["projectId"]
        }
    ),
    # Memory Tools (placeholders for Dev B to implement)
    Tool(
        name="complete_task",
        description="Mark a task as complete with documentation. Generates AI documentation and stores embeddings.",
        inputSchema={
            "type": "object",
            "properties": {
                "taskId": {"type": "string", "description": "ID of the task to complete"},
                "projectId": {"type": "string", "description": "ID of the project"},
                "userId": {"type": "string", "description": "ID of the user completing the task"},
                "whatIDid": {"type": "string", "description": "Description of what was done"},
                "codeSnippet": {"type": "string", "description": "Optional code snippet"}
            },
            "required": ["taskId", "projectId", "userId", "whatIDid"]
        }
    ),
    Tool(
        name="memory_search",
        description="Semantic search across project memory. Returns relevant log entries and AI-synthesized answer.",
        inputSchema={
            "type": "object",
            "properties": {
                "projectId": {"type": "string", "description": "ID of the project to search"},
                "query": {"type": "string", "description": "Search query"},
                "filters": {
                    "type": "object",
                    "properties": {
                        "userId": {"type": "string"},
                        "dateFrom": {"type": "string"},
                        "dateTo": {"type": "string"},
                        "tags": {"type": "array", "items": {"type": "string"}}
                    },
                    "description": "Optional filters"
                }
            },
            "required": ["projectId", "query"]
        }
    ),
    # Smart Query Tool - LLM-first natural language queries
    Tool(
        name="smart_query",
        description="Natural language query with context awareness. Understands 'yesterday', 'I', user names, task references. Use for questions like 'What did I do yesterday?' or 'How does auth work?'",
        inputSchema={
            "type": "object",
            "properties": {
                "projectId": {"type": "string", "description": "Project ID to query"},
                "query": {"type": "string", "description": "Natural language query"},
                "currentUserId": {"type": "string", "description": "ID of user making the query"},
                "currentDatetime": {"type": "string", "description": "Current datetime ISO format (optional)"}
            },
            "required": ["projectId", "query", "currentUserId"]
        }
    )
]


@server.list_tools()
async def handle_list_tools() -> list[Tool]:
    """Return all available tools."""
    return TOOLS


@server.call_tool()
async def handle_call_tool(name: str, arguments: dict) -> list[TextContent]:
    """Handle tool calls by dispatching to the appropriate function."""
    
    try:
        match name:
            # Project Tools
            case "create_project":
                result = create_project(
                    name=arguments["name"],
                    description=arguments["description"],
                    user_id=arguments["userId"]
                )
            case "list_projects":
                result = list_projects(user_id=arguments["userId"])
            case "join_project":
                result = join_project(
                    project_id=arguments["projectId"],
                    user_id=arguments["userId"]
                )
            
            # Task Tools
            case "create_task":
                result = create_task(
                    project_id=arguments["projectId"],
                    title=arguments["title"],
                    description=arguments.get("description"),
                    assigned_to=arguments.get("assignedTo")
                )
            case "list_tasks":
                result = list_tasks(
                    project_id=arguments["projectId"],
                    status=arguments.get("status")
                )
            case "list_activity":
                result = list_activity(
                    project_id=arguments["projectId"],
                    limit=arguments.get("limit", 20)
                )
            
            # Memory Tools (Dev B - implemented)
            case "complete_task":
                result = await complete_task(
                    task_id=arguments["taskId"],
                    project_id=arguments["projectId"],
                    user_id=arguments["userId"],
                    what_i_did=arguments["whatIDid"],
                    code_snippet=arguments.get("codeSnippet")
                )
            case "memory_search":
                result = await memory_search(
                    project_id=arguments["projectId"],
                    query=arguments["query"],
                    filters=arguments.get("filters")
                )

            # Smart Query Tool
            case "smart_query":
                result = await smart_query(
                    project_id=arguments["projectId"],
                    query=arguments["query"],
                    current_user_id=arguments["currentUserId"],
                    current_datetime=arguments.get("currentDatetime")
                )

            case _:
                result = {"error": f"Unknown tool: {name}"}
        
        return [TextContent(type="text", text=json.dumps(result, default=str))]
    
    except Exception as e:
        error_result = {"error": str(e)}
        return [TextContent(type="text", text=json.dumps(error_result))]


async def main():
    """Run the MCP server."""
    # Initialize database tables
    init_db()
    # Initialize vector store for embeddings
    init_vectorstore()
    print("Project Memory MCP Server starting...", flush=True)
    
    async with stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            server.create_initialization_options()
        )


if __name__ == "__main__":
    asyncio.run(main())