mattermost-dl/mm_search_mcp.py

#!/usr/bin/env python3
"""MCP server exposing mm_search.py search capabilities.

Prefer a single shared HTTP instance when multiple Cursor agents use this MCP
(stdio spawns one process per agent and they contend on the sqlite index).

HTTP (recommended for multi-agent):

    MM_MCP_TRANSPORT=http MM_MCP_HOST=127.0.0.1 MM_MCP_PORT=8766 \\
      /path/to/.venv/bin/python /path/to/mm_search_mcp.py

Cursor ~/.cursor/mcp.json:

    {
      "mcpServers": {
        "mattermost-search": {
          "url": "http://127.0.0.1:8766/mcp"
        }
      }
    }

Stdio (default, one process per client):

    {
      "mcpServers": {
        "mattermost-search": {
          "command": "/path/to/.venv/bin/python",
          "args": ["/path/to/mm_search_mcp.py"],
          "cwd": "/path/to/mattermost-dl",
          "env": {
            "MM_SEARCH_ROOT": "results",
            "MM_SEARCH_DB": "mattermost_index.sqlite"
          }
        }
      }
    }

Requires: pip install fastmcp
"""
from __future__ import annotations

import datetime as dt
import os
from typing import Any, Literal, Optional

from fastmcp import FastMCP

import mm_search as mm_search_lib
import mm_users

mcp = FastMCP("mattermost-search")

SearchMode = Literal["message", "file_name", "file_content"]


def _env_root() -> str:
    return os.environ.get("MM_SEARCH_ROOT", "results")


def _env_db() -> str:
    return os.environ.get("MM_SEARCH_DB", "mattermost_index.sqlite")


def _parse_date(value: str) -> dt.datetime:
    try:
        return mm_search_lib.parse_iso_date(value)
    except Exception as exc:
        raise ValueError(f"invalid date '{value}': {exc}") from exc


def _collect_matches(
    *,
    query: str,
    root: str,
    user: str,
    channel: str,
    team: str,
    after: dt.datetime,
    before: dt.datetime,
    regex: bool,
    search_mode: SearchMode,
    ignore_case: bool,
    limit: int = 10,
    db: str,
    no_index: bool,
) -> list[dict[str, Any]]:
    if search_mode != "message" and regex:
        raise ValueError("search_mode file_name/file_content do not support regex=True")

    effective_ignore_case = ignore_case if ignore_case is not None else (not regex)
    matcher = mm_search_lib.make_matcher(query, regex, effective_ignore_case)
    use_index = (not no_index) and os.path.exists(db)

    if search_mode in ("file_name", "file_content"):
        if not use_index:
            raise ValueError(
                "Attachment search (file_name/file_content) requires sqlite index. "
                "Create it with mm_index.py and do not set no_index=True."
            )
        mode = "name" if search_mode == "file_name" else "content"
        matches_iter = mm_search_lib.search_attachments_index(
            db_path=db,
            query=query,
            user=user,
            channel=channel,
            team=team,
            after=after,
            before=before,
            limit=limit,
            mode=mode,
        )
        output_limit = None
    elif use_index and not regex:
        matches_iter = mm_search_lib.search_index(
            db_path=db,
            query=query,
            user=user,
            channel=channel,
            team=team,
            after=after,
            before=before,
            limit=limit,
            ignore_case=effective_ignore_case,
        )
        output_limit = None
    else:
        files = list(mm_search_lib.iter_result_files(root))
        matches_iter = (
            m
            for path in files
            for m in mm_search_lib.search_file(
                path=path,
                matcher=matcher,
                user=user,
                channel=channel,
                team=team,
                after=after,
                before=before,
            )
        )
        output_limit = limit

    out: list[dict[str, Any]] = []
    count = 0
    for match in matches_iter:
        if output_limit is not None and count >= output_limit:
            break
        out.append(match)
        count += 1
    return out


def _format_human(matches: list[dict[str, Any]], show_file: bool) -> str:
    lines: list[str] = []
    for m in matches:
        p = m["post"]
        created = p.get("created", "")
        username = p.get("username", "")
        msg = (p.get("message") or "").replace("\n", "\\n")
        attachment = m.get("attachment")
        attachment_suffix = ""
        if isinstance(attachment, dict) and attachment.get("filename"):
            attachment_suffix = f" [file: {attachment['filename']}]"
        if show_file:
            file_rel = os.path.relpath(m["file"])
            idx = p.get("idx", "?")
            prefix = f"{file_rel}:{idx}: "
        else:
            prefix = ""
        channel = m.get("channel") or "?"
        team = m.get("team") or ""
        team_prefix = f"[{team}]" if team else ""
        lines.append(f"{prefix}{team_prefix}[{channel}] {created} {username}: {msg}{attachment_suffix}")
    return "\n".join(lines)


@mcp.tool
def mm_search(
    query: str,
    root: str = None,
    user: str = None,
    channel: str = None,
    team: str = None,
    after: str = None,
    before: str = None,
    regex: bool = False,
    search_mode: SearchMode = "message",
    ignore_case: bool = None,
    limit: int = 10,
    db: str = None,
    no_index: bool = False,
    output: Literal["json", "human"] = "json",
    show_file: bool = False,
) -> dict[str, Any]:
    """Search Mattermost export JSON files produced by mattermost-dl.py.

    Mirrors mm_search.py CLI. Returns structured matches or human-readable lines.

    Args:
        query: Search string or regex (when regex=True).
        root: Root directory with export results (default: MM_SEARCH_ROOT or "results").
        user: Filter by username (exact match).
        channel: Filter by channel display_name or name (exact match).
        team: Filter by team name (exact match).
        after: Only messages on/after this date (YYYY-MM-DD or ISO).
        before: Only messages on/before this date (YYYY-MM-DD or ISO).
        regex: Treat query as regular expression (message scan only).
        search_mode: "message" (default), "file_name", or "file_content".
        ignore_case: Case-insensitive search. Default: True for substring, False for regex.
        limit: Maximum number of matches to return (default: 10).
        db: SQLite index path (default: MM_SEARCH_DB or "mattermost_index.sqlite").
        no_index: Do not use sqlite index even if DB exists.
        output: "json" returns match objects; "human" returns CLI-style text lines.
        show_file: Include source filename and per-file index in human output.
    """
    resolved_root = root or _env_root()
    resolved_db = db or _env_db()
    parsed_after = _parse_date(after) if after else None
    parsed_before = _parse_date(before) if before else None

    matches = _collect_matches(
        query=query,
        root=resolved_root,
        user=user,
        channel=channel,
        team=team,
        after=parsed_after,
        before=parsed_before,
        regex=regex,
        search_mode=search_mode,
        ignore_case=ignore_case,
        limit=limit,
        db=resolved_db,
        no_index=no_index,
    )

    result: dict[str, Any] = {
        "count": len(matches),
        "used_index": (not no_index)
        and os.path.exists(resolved_db)
        and (search_mode in ("file_name", "file_content") or not regex),
        "search_mode": search_mode,
        "root": resolved_root,
        "db": resolved_db,
    }

    if output == "human":
        result["text"] = _format_human(matches, show_file=show_file)
    else:
        result["matches"] = matches

    return result


@mcp.tool
def mm_resolve_user(
    ref: str,
    db: str = None,
) -> dict[str, Any]:
    """Resolve display name, FIO, nickname, or username to Mattermost username.

    Args:
        ref: User reference — username, first/last name, nickname, or FIO
             (e.g. "Цисык", "vladislav.tsisyk").
        db: SQLite index path (default: MM_SEARCH_DB or "mattermost_index.sqlite").
    """
    resolved_db = db or _env_db()
    return mm_users.resolve_user(resolved_db, ref)


@mcp.tool
def mm_search_users(
    query: str,
    limit: int = 10,
    db: str = None,
) -> dict[str, Any]:
    """Search users in the local directory by name or username.

    Args:
        query: Substring to match against username, first/last name, nickname, full name.
        limit: Maximum matches (default: 10).
        db: SQLite index path (default: MM_SEARCH_DB or "mattermost_index.sqlite").
    """
    resolved_db = db or _env_db()
    matches = mm_users.search_users(resolved_db, query, limit=limit)
    return {
        "query": query,
        "count": len(matches),
        "db": resolved_db,
        "matches": matches,
    }


@mcp.tool
def mm_user_directory_refresh(
    root: str = None,
    db: str = None,
    from_api: bool = False,
    config_path: str = None,
) -> dict[str, Any]:
    """Rebuild the local user directory from users.json, Mattermost API, and post usernames.

    Args:
        root: Export root (default: MM_SEARCH_ROOT or "results").
        db: SQLite index path (default: MM_SEARCH_DB).
        from_api: Also fetch profiles from Mattermost API (needs config.json with token).
        config_path: Path to Mattermost config (default: MM_CONFIG_PATH or "config.json").
    """
    resolved_root = root or _env_root()
    resolved_db = db or _env_db()
    return mm_users.refresh_user_directory(
        resolved_db,
        resolved_root,
        config_path=config_path,
        from_api=from_api,
    )


@mcp.tool
def mm_search_tag(
    keyword: str,
    after: str,
    before: str,
    user: str = None,
    channel: str = None,
    team: str = None,
    limit: int = 50,
    db: str = None,
) -> dict[str, Any]:
    """Find posts containing #keyword in a required date range (newest first).

    Plain SQL substring match — no FTS/rank. Use for daily tags like plan/report
    when you have after/before (e.g. one day or one week).

    Args:
        keyword: Tag without or with leading # (e.g. "plan", "report", "#plan").
        after: Inclusive start date (YYYY-MM-DD or ISO) — required.
        before: Inclusive end date (YYYY-MM-DD or ISO) — required.
        user: Filter by username (exact match).
        channel: Filter by channel display_name or name (exact match).
        team: Filter by team name (exact match).
        limit: Max matches (default: 50).
        db: SQLite index path (default: MM_SEARCH_DB).
    """
    resolved_db = db or _env_db()
    parsed_after = _parse_date(after)
    parsed_before = _parse_date(before)
    # Date-only before → end of that UTC day so a single day is inclusive.
    if "T" not in before.strip():
        parsed_before = parsed_before + dt.timedelta(days=1) - dt.timedelta(microseconds=1)

    matches = list(
        mm_search_lib.search_tag(
            resolved_db,
            keyword,
            parsed_after,
            parsed_before,
            user=user,
            channel=channel,
            team=team,
            limit=limit,
        )
    )
    tag = keyword.strip()
    if not tag.startswith("#"):
        tag = f"#{tag}"
    return {
        "count": len(matches),
        "used_index": True,
        "search_mode": "tag",
        "tag": tag,
        "after": after,
        "before": before,
        "db": resolved_db,
        "matches": matches,
    }


@mcp.tool
def mm_search_index_status(
    db: Optional[str] = None,
    root: Optional[str] = None,
) -> dict[str, Any]:
    """Report whether the sqlite index and export root are available."""
    resolved_db = db or _env_db()
    resolved_root = root or _env_root()
    db_exists = os.path.exists(resolved_db)
    root_exists = os.path.isdir(resolved_root)

    json_count = 0
    if root_exists:
        json_count = sum(1 for _ in mm_search_lib.iter_result_files(resolved_root))

    user_count = 0
    if db_exists:
        try:
            conn = mm_search_lib.open_db(resolved_db, readonly=True)
            try:
                user_count = conn.execute("SELECT COUNT(*) FROM users").fetchone()[0]
            finally:
                conn.close()
        except Exception:
            user_count = None

    return {
        "db": resolved_db,
        "db_exists": db_exists,
        "root": resolved_root,
        "root_exists": root_exists,
        "json_files": json_count,
        "users": user_count,
        "index_required_for": ["file_name", "file_content"],
    }


if __name__ == "__main__":
    transport = os.environ.get("MM_MCP_TRANSPORT", "stdio").strip().lower()
    if transport in ("http", "streamable-http", "sse"):
        host = os.environ.get("MM_MCP_HOST", "127.0.0.1")
        port = int(os.environ.get("MM_MCP_PORT", "8766"))
        path = os.environ.get("MM_MCP_PATH", "/mcp")
        # Stateless HTTP: safe for multiple Cursor agents against one process.
        mcp.run(
            transport="http" if transport == "streamable-http" else transport,
            host=host,
            port=port,
            path=path,
            stateless_http=True,
        )
    else:
        mcp.run(transport="stdio")