38775-vm/core/mcp.py

import json
import os
import httpx
from django.conf import settings
from .models import MCPToolRequest, HumanResponse, SlackSettings
from asgiref.sync import sync_to_async

SLACK_API_URL = "https://slack.com/api/chat.postMessage"

@sync_to_async
def get_slack_token():
    token = os.getenv("SLACK_BOT_TOKEN")
    if token:
        return token
    config = SlackSettings.objects.first()
    if config and config.bot_token:
        return config.bot_token
    return None

async def send_slack_message(channel, text, thread_ts=None):
    token = await get_slack_token()
    if not token:
        return {"error": "SLACK_BOT_TOKEN not configured", "ok": False}
    
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json; charset=utf-8"
    }
    payload = {
        "channel": channel,
        "text": text,
    }
    if thread_ts:
        payload["thread_ts"] = thread_ts

    async with httpx.AsyncClient() as client:
        response = await client.post(SLACK_API_URL, headers=headers, json=payload)
        data = response.json()
        return data

async def handle_mcp_request(body):
    """
    Simple MCP JSON-RPC handler for HTTP/SSE.
    Supports tools/list and tools/call.
    """
    try:
        request_data = json.loads(body)
    except json.JSONDecodeError:
        return {"jsonrpc": "2.0", "error": {"code": -32700, "message": "Parse error"}, "id": None}

    method = request_data.get("method")
    params = request_data.get("params", {})
    request_id = request_data.get("id")

    if method == "initialize":
        return {
            "jsonrpc": "2.0",
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {}
                },
                "serverInfo": {
                    "name": "flatlogic-slack",
                    "version": "1.0.0"
                }
            },
            "id": request_id
        }

    elif method == "notifications/initialized":
        return None  # No response for notifications

    elif method == "ping":
        return {
            "jsonrpc": "2.0",
            "result": {},
            "id": request_id
        }

    elif method == "tools/list":
        return {
            "jsonrpc": "2.0",
            "result": {
                "tools": [
                    {
                        "name": "send_slack_message",
                        "description": "Sends a message to a Slack channel.",
                        "inputSchema": {
                            "type": "object",
                            "properties": {
                                "channel": {"type": "string", "description": "Channel ID or Name"},
                                "text": {"type": "string", "description": "Message content"}
                            },
                            "required": ["channel", "text"]
                        }
                    },
                    {
                        "name": "list_responses",
                        "description": "Lists human responses for a specific request ID.",
                        "inputSchema": {
                            "type": "object",
                            "properties": {
                                "request_id": {"type": "string", "description": "The UUID of the original request"}
                            },
                            "required": ["request_id"]
                        }
                    }
                ]
            },
            "id": request_id
        }

    elif method == "tools/call":
        tool_name = params.get("name")
        arguments = params.get("arguments", {})

        if tool_name == "send_slack_message":
            channel = arguments.get("channel")
            text = arguments.get("text")

            # Log request
            db_request = await MCPToolRequest.objects.acreate(
                tool_name=tool_name,
                arguments=arguments,
                slack_channel=channel or ""
            )

            slack_result = await send_slack_message(channel, text)
            
            if slack_result.get("ok"):
                db_request.status = 'SENT'
                db_request.slack_ts = slack_result.get("ts", "")
                db_request.response_json = slack_result
                await db_request.asave()
                
                return {
                    "jsonrpc": "2.0",
                    "result": {
                        "content": [
                            {
                                "type": "text",
                                "text": f"Message sent to Slack (TS: {db_request.slack_ts}). Request ID: {db_request.id}"
                            }
                        ]
                    },
                    "id": request_id
                }
            else:
                db_request.status = 'ERROR'
                db_request.response_json = slack_result
                await db_request.asave()
                return {
                    "jsonrpc": "2.0",
                    "error": {
                        "code": -32000,
                        "message": f"Slack API Error: {slack_result.get('error')}"
                    },
                    "id": request_id
                }

        elif tool_name == "list_responses":
            req_id = arguments.get("request_id")
            responses = []
            async for r in HumanResponse.objects.filter(request_id=req_id):
                responses.append({
                    "text": r.text,
                    "user": r.user_name,
                    "file_url": r.file_url,
                    "timestamp": str(r.created_at)
                })
            
            return {
                "jsonrpc": "2.0",
                "result": {
                    "content": [
                        {
                            "type": "text",
                            "text": json.dumps(responses, indent=2) if responses else "No responses yet."
                        }
                    ]
                },
                "id": request_id
            }

    return {
        "jsonrpc": "2.0",
        "error": {"code": -32601, "message": "Method not found"},
        "id": request_id
    }