vortex/mcp/server/vortex_mcp/server.py

"""Vortex MCP server: read-first Matrix access with a human-gated draft outbox."""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
import sqlite3
import uuid
from collections import deque
from datetime import datetime, timezone
from pathlib import Path

from fastmcp import FastMCP
from fastmcp.exceptions import ToolError
from nio import (
    AsyncClient,
    AsyncClientConfig,
    LoginResponse,
    MessageDirection,
    RoomMessagesResponse,
    RoomMessageText,
)
from nio.crypto import ENCRYPTION_ENABLED

log = logging.getLogger("vortex-mcp")

CREDENTIALS_FILE = os.environ.get("CREDENTIALS_FILE", "/shared/credentials.json")
STORE_DIR = Path(os.environ.get("STORE_DIR", "/shared/mcp-store"))
DRAFTS_DB = os.environ.get("DRAFTS_DB", "/shared/drafts/drafts.db")
NETWORKS = ("gmessages", "telegram", "whatsapp", "twitter", "linkedin", "discord", "meta")

mcp = FastMCP("vortex")

_client: AsyncClient | None = None
_sync_task: asyncio.Task | None = None
_init_lock = asyncio.Lock()
# room_id -> deque of message dicts, filled by the background sync loop
_cache: dict[str, deque] = {}


# ---------------------------------------------------------------- drafts db

def _db() -> sqlite3.Connection:
    Path(DRAFTS_DB).parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(DRAFTS_DB)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute(
        """CREATE TABLE IF NOT EXISTS drafts(
             draft_id TEXT PRIMARY KEY, room_id TEXT NOT NULL, body TEXT NOT NULL,
             reply_to_event_id TEXT, created_by TEXT NOT NULL, created_at TEXT NOT NULL,
             status TEXT NOT NULL DEFAULT 'draft', sent_event_id TEXT)"""
    )
    return conn


def _get_draft(conn: sqlite3.Connection, draft_id: str) -> dict:
    row = conn.execute("SELECT * FROM drafts WHERE draft_id=?", (draft_id,)).fetchone()
    if row is None:
        raise ToolError(f"Draft not found: {draft_id}")
    return dict(row)


# ------------------------------------------------------------- matrix client

async def _load_credentials() -> dict:
    for _ in range(150):  # bootstrap provisions the file; wait up to ~5 min
        if os.path.exists(CREDENTIALS_FILE):
            with open(CREDENTIALS_FILE) as f:
                return json.load(f)
        log.info("Waiting for credentials file %s ...", CREDENTIALS_FILE)
        await asyncio.sleep(2)
    raise ToolError(f"Credentials file never appeared: {CREDENTIALS_FILE}")


async def _login(creds: dict) -> AsyncClient:
    homeserver = os.environ.get("MATRIX_HOMESERVER") or creds["homeserver"]
    STORE_DIR.mkdir(parents=True, exist_ok=True)
    session_file = STORE_DIR / "session.json"
    config = AsyncClientConfig(
        encryption_enabled=ENCRYPTION_ENABLED, store_sync_tokens=True
    )
    client = AsyncClient(
        homeserver, creds["user_id"], store_path=str(STORE_DIR), config=config
    )

    if session_file.exists():
        sess = json.loads(session_file.read_text())
        client.restore_login(sess["user_id"], sess["device_id"], sess["access_token"])
        log.info("Restored session for %s (device %s)", sess["user_id"], sess["device_id"])
        return client

    resp = await client.login(creds.get("password") or "", device_name="vortex-mcp")
    if isinstance(resp, LoginResponse):
        session_file.write_text(json.dumps({
            "user_id": resp.user_id,
            "device_id": resp.device_id,
            "access_token": resp.access_token,
        }))
        log.info("Password login ok, session persisted")
        return client

    # ponytail: fall back to the shared bootstrap token if password login fails
    log.warning("Password login failed (%s); using shared access_token", resp)
    client.restore_login(creds["user_id"], creds.get("device_id") or "", creds["access_token"])
    return client


async def _on_message(room, event) -> None:
    dq = _cache.setdefault(room.room_id, deque(maxlen=1000))
    dq.append({
        "event_id": event.event_id,
        "sender": event.sender,
        "sender_display": room.user_name(event.sender) or event.sender,
        "body": event.body,
        "timestamp_ms": event.server_timestamp,
        "room_id": room.room_id,
    })


async def _get_client() -> AsyncClient:
    """Lazy-init the nio client and background sync task on first tool call."""
    global _client, _sync_task
    async with _init_lock:
        if _client is None:
            creds = await _load_credentials()
            _client = await _login(creds)
            _client.add_event_callback(_on_message, RoomMessageText)
        if _sync_task is None or _sync_task.done():
            _sync_task = asyncio.create_task(
                _client.sync_forever(timeout=30000, full_state=True)
            )
    return _client


def _room_network(client: AsyncClient, room) -> str:
    domain = client.user_id.split(":", 1)[1] if ":" in (client.user_id or "") else ""
    for uid in room.users:
        for net in NETWORKS:
            if uid.startswith(f"@{net}_") or uid == f"@{net}bot:{domain}":
                return net
    return "matrix"


# -------------------------------------------------------------------- tools

@mcp.tool
async def list_rooms(network_filter: str | None = None) -> list[dict]:
    """List joined rooms. Optionally filter by bridge network
    (gmessages, telegram, whatsapp, twitter, linkedin, discord, meta, matrix)."""
    client = await _get_client()
    out = []
    for room_id, room in client.rooms.items():
        network = _room_network(client, room)
        if network_filter and network != network_filter:
            continue
        out.append({
            "room_id": room_id,
            "name": room.display_name,
            "network": network,
            "unread_count": getattr(room, "unread_notifications", 0) or 0,
            "member_count": room.member_count,
        })
    return out


@mcp.tool
async def read_messages(room_id: str, limit: int = 50, before: str | None = None) -> dict:
    """Read recent messages from a room, newest first. Pass the returned `end`
    token as `before` to paginate further back."""
    client = await _get_client()
    resp = await client.room_messages(
        room_id,
        start=before or client.next_batch,
        direction=MessageDirection.back,
        limit=limit,
    )
    if not isinstance(resp, RoomMessagesResponse):
        raise ToolError(f"room_messages failed: {resp}")
    room = client.rooms.get(room_id)
    messages = [
        {
            "event_id": ev.event_id,
            "sender": ev.sender,
            "sender_display": room.user_name(ev.sender) if room else ev.sender,
            "body": ev.body,
            "timestamp_ms": ev.server_timestamp,
            "room_id": room_id,
        }
        for ev in resp.chunk
        if isinstance(ev, RoomMessageText)
    ]
    return {"messages": messages, "end": resp.end}


@mcp.tool
async def search_messages(query: str, room_id: str | None = None) -> list[dict]:
    """Case-insensitive substring search over messages seen since this server
    started syncing. Optionally restrict to one room."""
    await _get_client()
    needle = query.lower()
    rooms = [room_id] if room_id else list(_cache)
    return [
        msg
        for rid in rooms
        for msg in _cache.get(rid, ())
        if needle in msg["body"].lower()
    ]


@mcp.tool
def create_draft(room_id: str, body: str, reply_to_event_id: str | None = None) -> dict:
    """Create a message draft for a room. Never sends anything; a human reviews
    drafts and sends them from the TUI."""
    draft = {
        "draft_id": uuid.uuid4().hex,
        "room_id": room_id,
        "body": body,
        "reply_to_event_id": reply_to_event_id,
        "created_by": "mcp",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "status": "draft",
        "sent_event_id": None,
    }
    with _db() as conn:
        conn.execute(
            "INSERT INTO drafts VALUES(:draft_id,:room_id,:body,:reply_to_event_id,"
            ":created_by,:created_at,:status,:sent_event_id)",
            draft,
        )
    return draft


@mcp.tool
def list_drafts() -> list[dict]:
    """List all pending (unsent, undiscarded) drafts."""
    with _db() as conn:
        rows = conn.execute("SELECT * FROM drafts WHERE status='draft'").fetchall()
    return [dict(r) for r in rows]


@mcp.tool
def update_draft(draft_id: str, body: str) -> dict:
    """Replace the body of a pending draft."""
    with _db() as conn:
        draft = _get_draft(conn, draft_id)
        if draft["status"] != "draft":
            raise ToolError(f"Draft {draft_id} is {draft['status']}, not editable")
        conn.execute("UPDATE drafts SET body=? WHERE draft_id=?", (body, draft_id))
        draft["body"] = body
    return draft


@mcp.tool
def discard_draft(draft_id: str) -> dict:
    """Mark a draft as discarded."""
    with _db() as conn:
        draft = _get_draft(conn, draft_id)
        conn.execute("UPDATE drafts SET status='discarded' WHERE draft_id=?", (draft_id,))
        draft["status"] = "discarded"
    return draft


@mcp.tool
async def send_draft(draft_id: str) -> dict:
    """Send a pending draft to its room. Only works when MCP_ALLOW_SEND=true;
    otherwise a human must send it from the TUI."""
    if os.environ.get("MCP_ALLOW_SEND", "false").lower() != "true":
        raise ToolError(
            "Sending disabled: a human must send this draft from the TUI (MCP_ALLOW_SEND=false)"
        )
    with _db() as conn:
        draft = _get_draft(conn, draft_id)
    if draft["status"] != "draft":
        raise ToolError(f"Draft {draft_id} is {draft['status']}, cannot send")

    client = await _get_client()
    content: dict = {"msgtype": "m.text", "body": draft["body"]}
    if draft["reply_to_event_id"]:
        content["m.relates_to"] = {
            "m.in_reply_to": {"event_id": draft["reply_to_event_id"]}
        }
    resp = await client.room_send(
        draft["room_id"], "m.room.message", content, ignore_unverified_devices=True
    )
    event_id = getattr(resp, "event_id", None)
    if not event_id:
        raise ToolError(f"Send failed: {resp}")
    with _db() as conn:
        conn.execute(
            "UPDATE drafts SET status='sent', sent_event_id=? WHERE draft_id=?",
            (event_id, draft_id),
        )
    draft.update(status="sent", sent_event_id=event_id)
    return draft


# --------------------------------------------------------------------- main

def main() -> None:
    logging.basicConfig(level=logging.INFO)
    parser = argparse.ArgumentParser(description="Vortex MCP server")
    parser.add_argument("--stdio", action="store_true", help="use stdio transport")
    args = parser.parse_args()
    if args.stdio:
        mcp.run()
    else:
        mcp.run(
            transport="http",
            host="0.0.0.0",
            port=int(os.environ.get("MCP_HTTP_PORT", "8765")),
        )


if __name__ == "__main__":
    main()